tf.data 的预取优化不起作用

Question

我正在使用 tf.data API 并分析通过 here 编写的优化获得的各种加速。 我正在使用 tf.data API，我正在分析通过写入的优化获得的各种加速，但在所有情况下，我注意到您没有使用预取选项优化性能。 似乎没有实施任何优化，因此 CPU 和 GPU 之间没有重叠。 目前我使用的是 TF 2.11.0，但我也使用过 TF 2.10.0 和 TF 2.8.3，事实是一样的。 我也尝试过不同的批量大小。 我也使用过不同的 PC 和不同的 GPU，但事实是一样的。 我使用 Cifar10，其中每张图像都是 RGB 32x32，这个训练集由 40.000 张图像组成。 我使用的虚拟代码是这样的：

import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.layers as layers
import time

def get_model_data_augmentation_CPU():
    """Return the Keras model for data-augmentation on CPU"""
    # Define Keras Model
    model = tf.keras.Sequential([
      layers.Conv2D(64, 3, activation='relu'),
      layers.MaxPooling2D(),
      layers.Dropout(0.1),
      layers.Conv2D(128, 3, activation='relu'),
      layers.MaxPooling2D(),
      layers.Dropout(0.1),
      layers.Conv2D(128, 3, activation='relu'),
      layers.MaxPooling2D(),
      layers.Dropout(0.2),
      layers.Flatten(),
      layers.Dense(256, activation='relu'),
      layers.Dropout(0.3),
      layers.Dense(10)
    ])
    adam_opt = keras.optimizers.Adam(learning_rate=0.001)
    model.compile(optimizer = adam_opt,
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])
    return model

data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip(mode='horizontal'),
    tf.keras.layers.RandomRotation(0.1),
    tf.keras.layers.RandomZoom(0.2),
])

model = get_model_data_augmentation_CPU()

BATCH_SIZE = 32 
(X_train, y_train), (X_test, y_test) = keras.datasets.cifar10.load_data()
dataset_train = tf.data.Dataset.from_tensor_slices((X_train, y_train)) 
dataset_train = dataset_train.map(lambda x,y : (data_augmentation(x),y), num_parallel_calls=3) 
dataset_train = dataset_train.batch(BATCH_SIZE)
dataset_train = dataset_train.prefetch(1) # If I comment this line the perfomance remain the same


start_time = time.time()
history = model.fit(
    dataset_train,
    epochs=EPOCHS,
)
end_time = time.time()

Answer 1

好的，我得到了结果。
我创建了一个笔记本并在这个数据集上进行了测试：
https://www.kaggle.com/competitions/dogs-vs-cats/data

我用预取运行笔记本 3 次，不使用预取运行 3 次，结果如下：
有预取，训练时间：43.15s，44.12s，43.53s
没有预取，训练时间：45.97s，46.62s，45.95s
这基本上是 5.6% 的改进。还不错。

我认为在你的情况下没有太大的改进，因为你所有的数据集都已经在你的 ram 中并且处理不需要太多时间，所以基本上没有区别。

如果您想尝试一下，这是代码（它基于您的代码）：

import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.layers as layers
import time

def get_model():
    """Return the Keras model for data-augmentation on CPU"""
    # Define Keras Model
    model = tf.keras.Sequential([
      layers.Conv2D(64, 3, activation='relu'),
      layers.MaxPooling2D(),
      layers.Dropout(0.1),
      layers.Conv2D(128, 3, activation='relu'),
      layers.MaxPooling2D(),
      layers.Dropout(0.1),
      layers.Conv2D(128, 3, activation='relu'),
      layers.MaxPooling2D(),
      layers.Dropout(0.2),
      layers.Flatten(),
      layers.Dense(256, activation='relu'),
      layers.Dropout(0.3),
      layers.Dense(10, activation="softmax")
    ])
    return model

from tensorflow.keras.utils import image_dataset_from_directory
data_dir = "catdogs"
train_dataset = keras.preprocessing.image_dataset_from_directory(
    data_dir,
    seed=123,
    image_size=(128,128),
    batch_size=32,
    subset="training", #subset="both", lo nuevo
    shuffle=True,
    validation_split=0.1,
) 
val_dataset = keras.preprocessing.image_dataset_from_directory(
    data_dir,
    seed=123,
    image_size=(128,128),
    batch_size=32,
    subset="validation", 
    shuffle=False,
    validation_split=0.1,
) 
data_augmentation = tf.keras.Sequential([
    layers.Rescaling(1./255),
    layers.RandomRotation(0.2),
    layers.RandomFlip(mode='horizontal'),
    layers.RandomZoom(0.2),
])
model = keras.Sequential([
    data_augmentation,
    get_model()
])
adam_opt = keras.optimizers.Adam(learning_rate=0.001)
model.compile(optimizer = adam_opt,
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

start_time = time.time()
history = model.fit(
    train_dataset,
    epochs=2,
    validation_data=val_dataset
)
end_time = time.time()
print(f"Total time: {end_time - start_time:.2f}")