這個簡單的範例示範如何將 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：建立您的輸入管線
首先，根據以下建議建立有效率的輸入管線
載入資料集
使用以下引數載入 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：在隨機排序後批次處理資料集元素，以在每個週期獲得獨特的批次。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。 - 快取會在批次處理後完成，因為批次在週期之間可能相同。
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>