如何在Google Colab训练CNN前预加载PlantVillage数据集至GPU？

1. 提前完成数据预加载与缓存 #

用TensorFlow的数据集工具把数据提前加载并缓存，避免训练时首次加载拖慢速度：

• 先加载数据集并搭建预处理流水线（这部分不计时）：

  import tensorflow as tf
  from tensorflow.keras.utils import image_dataset_from_directory
  
  # 从云端硬盘加载PlantVillage子集
  train_ds = image_dataset_from_directory(
      '/content/drive/MyDrive/PlantVillage_subset',
      image_size=(224, 224),
      batch_size=32
  )
  
  # 添加缓存和预取，让数据加载流水线更高效
  train_ds = train_ds.cache()  # 小数据集直接缓存到内存，大数据集可指定磁盘路径
  train_ds = train_ds.prefetch(tf.data.AUTOTUNE)
  

• 手动遍历一次数据集，触发首次数据加载（这部分也不计时）：

  # 跑一遍数据集，把数据提前加载到缓存/GPU，后续训练直接用缓存
  for _ in train_ds:
      pass
  

2. 单独给训练代码计时 #

完成上面的预加载后，只把模型训练的代码用%%time包裹，这样计时就不会包含数据加载的耗时：

%%time
# 这里放你的模型定义和训练代码
model = tf.keras.Sequential([
    # 你的CNN层定义
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_ds, epochs=10)

额外优化提示 #

如果数据集太大内存装不下，把cache()改成cache('/content/temp_cache')，用磁盘缓存代替内存缓存，虽然速度比内存稍慢，但远快于每次从云端硬盘读取。prefetch能让CPU在GPU处理当前批次时提前准备下一批数据，进一步提升训练效率。

内容的提问来源于stack exchange，提问作者Maral