这个简单的示例演示了如何将 TensorFlow Datasets (TFDS) 插入 Keras 模型。
在 TensorFlow.org 上查看 | 在 Google Colab 中运行 | 在 GitHub 上查看源代码 | 下载笔记本 |
import tensorflow as tf
import tensorflow_datasets as tfds
2023-10-03 09:29:30.258272: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2023-10-03 09:29:30.258321: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2023-10-03 09:29:30.258358: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
步骤 1:创建您的输入管道
首先使用以下建议构建一个高效的输入管道
- 性能提示 指南
- 使用
tf.data
API 提高性能 指南
加载数据集
使用以下参数加载 MNIST 数据集
shuffle_files=True
:MNIST 数据仅存储在一个文件中,但对于磁盘上有多个文件的更大数据集,在训练时对它们进行混洗是一个好习惯。as_supervised=True
:返回一个元组(img, label)
而不是一个字典{'image': img, 'label': label}
。
(ds_train, ds_test), ds_info = tfds.load(
'mnist',
split=['train', 'test'],
shuffle_files=True,
as_supervised=True,
with_info=True,
)
2023-10-03 09:29:33.682941: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:268] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
构建训练管道
应用以下转换
tf.data.Dataset.map
:TFDS 提供类型为tf.uint8
的图像,而模型期望类型为tf.float32
的图像。因此,您需要对图像进行归一化。tf.data.Dataset.cache
由于您将数据集放入内存中,因此在混洗之前将其缓存以提高性能。
注意: 随机转换应在缓存之后应用。tf.data.Dataset.shuffle
:为了实现真正的随机性,将混洗缓冲区设置为完整的数据集大小。
注意: 对于无法放入内存的大型数据集,如果您的系统允许,请使用buffer_size=1000
。tf.data.Dataset.batch
:在混洗之后对数据集的元素进行批处理,以便在每个 epoch 中获得唯一的批次。tf.data.Dataset.prefetch
:最好通过预取 来结束管道以提高性能。
def normalize_img(image, label):
"""Normalizes images: `uint8` -> `float32`."""
return tf.cast(image, tf.float32) / 255., label
ds_train = ds_train.map(
normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(128)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)
构建评估管道
您的测试管道与训练管道类似,但有一些细微差别
- 您不需要调用
tf.data.Dataset.shuffle
。 - 缓存是在批处理之后完成的,因为批次在 epoch 之间可以相同。
ds_test = ds_test.map(
normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_test = ds_test.batch(128)
ds_test = ds_test.cache()
ds_test = ds_test.prefetch(tf.data.AUTOTUNE)
步骤 2:创建和训练模型
将 TFDS 输入管道插入一个简单的 Keras 模型,编译模型并对其进行训练。
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10)
])
model.compile(
optimizer=tf.keras.optimizers.Adam(0.001),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
)
model.fit(
ds_train,
epochs=6,
validation_data=ds_test,
)
Epoch 1/6 469/469 [==============================] - 4s 4ms/step - loss: 0.3621 - sparse_categorical_accuracy: 0.9011 - val_loss: 0.1925 - val_sparse_categorical_accuracy: 0.9463 Epoch 2/6 469/469 [==============================] - 1s 3ms/step - loss: 0.1602 - sparse_categorical_accuracy: 0.9543 - val_loss: 0.1392 - val_sparse_categorical_accuracy: 0.9588 Epoch 3/6 469/469 [==============================] - 1s 2ms/step - loss: 0.1174 - sparse_categorical_accuracy: 0.9664 - val_loss: 0.1084 - val_sparse_categorical_accuracy: 0.9693 Epoch 4/6 469/469 [==============================] - 1s 3ms/step - loss: 0.0911 - sparse_categorical_accuracy: 0.9743 - val_loss: 0.0968 - val_sparse_categorical_accuracy: 0.9714 Epoch 5/6 469/469 [==============================] - 1s 2ms/step - loss: 0.0738 - sparse_categorical_accuracy: 0.9790 - val_loss: 0.0881 - val_sparse_categorical_accuracy: 0.9735 Epoch 6/6 469/469 [==============================] - 1s 2ms/step - loss: 0.0617 - sparse_categorical_accuracy: 0.9823 - val_loss: 0.0793 - val_sparse_categorical_accuracy: 0.9749 <keras.src.callbacks.History at 0x7fc41e0cb880>