谷歌Colab训练植物病害识别CNN模型卡在第一个Epoch如何解决

• Google Drive IO瓶颈：当前数据集存储在挂载的Google Drive中，Google Drive的网络读写速度远低于Colab运行时本地磁盘，代码中使用.cache()方法需要在首个Epoch启动前把所有训练数据加载到内存，从Drive读取大量图片的过程会表现为训练卡在第一个Epoch无进展。
• 内存不足：如果数据集规模过大，.cache()尝试将全量数据存入内存时超出Colab免费实例的内存上限，会导致进程阻塞无报错。
• 配置参数不匹配：batch size设置过大超出Colab GPU显存容量、shuffle_size设置超出数据集总样本数等情况，也可能导致训练启动阶段卡住。

1. 优先解决IO问题：将数据集从Google Drive复制到Colab本地磁盘后再读取，示例操作代码：

# 复制数据集压缩包到本地（如果是未压缩的文件夹，使用 cp -r 命令）
!cp /content/drive/MyDrive/google-colab-files/PlantVillage.zip /content/
# 解压数据集
!unzip -q /content/PlantVillage.zip -d /content/

之后将image_dataset_from_directory的路径修改为本地路径/content/PlantVillage即可。
2. 调整缓存策略：

• 若数据集规模较小，复制到本地后保留.cache()配置即可正常运行
• 若数据集规模较大，删除.cache()方法，仅保留.prefetch(buffer_size = tf.data.AUTOTUNE)即可，避免内存不足

3. 调整训练参数：

• 逐步降低batch size到32、16，排查是否是显存不足导致的阻塞
• 将shuffle_size调整为不超过数据集总样本数的数值，避免shuffle阶段异常等待

4. 提前验证数据加载：在启动训练前先执行print(len(train_ds))，如果该命令可以正常返回训练集的批次数，说明数据加载链路正常，再启动训练即可。

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