如何在Google Colab上高效训练含110k+图像的PyTorch CNN模型？

一、优化H5数据加载效率 #

• 持久化H5文件句柄：不要在__getitem__里反复打开H5文件，在Dataset的__init__中一次性打开文件并保持句柄，用内存映射模式（比如h5py.File('data.h5', 'r', swmr=True)），避免重复IO开销。同时把H5内的数据集转为numpy内存视图，减少每次索引的耗时。
• 调优DataLoader参数：设置num_workers=4-8（根据Colab实例CPU核心数调整），搭配persistent_workers=True，让数据加载子进程在epoch间不销毁，减少重复初始化的时间。开启pin_memory=True，直接把数据锁进GPU内存，加快数据传输速度。
• 重构H5文件结构：如果当前H5是按单张图像存储的小条目，改成批量块存储（比如每1000张存为一个数据集块），读取时批量加载，降低随机IO的频次。

二、最大化硬件利用率 #

• 开启混合精度训练：用PyTorch的torch.cuda.amp模块，在不损失精度的前提下减少显存占用、提升计算速度，只需几行代码改造：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  for inputs, labels in dataloader:
      optimizer.zero_grad()
      with autocast():
          outputs = model(inputs)
          loss = criterion(outputs, labels)
      scaler.scale(loss).backward()
      scaler.step(optimizer)
      scaler.update()
  

• 切换高级GPU实例：通过Colab的「运行时」->「更改运行时类型」选择A100等高端GPU，这类实例的显存和计算能力远高于普通T4，能支持更大的batch size，直接压缩单轮训练时间。
• 释放闲置资源：关闭Colab页面外的闲置标签、停止无关的后台运行代码，避免占用CPU和内存影响数据加载效率。

三、调整训练策略 #

• 放大batch size（或梯度累积）：在显存允许范围内尽量调大batch size（比如从32升到64/128）；如果显存不足，用梯度累积模拟大batch效果——每N次迭代再更新一次参数：

  accumulation_steps = 4
  for i, (inputs, labels) in enumerate(dataloader):
      optimizer.zero_grad()
      outputs = model(inputs)
      loss = criterion(outputs, labels)
      loss = loss / accumulation_steps
      loss.backward()
      if (i + 1) % accumulation_steps == 0:
          optimizer.step()
  

• 减少冗余计算：训练时关闭模型的dropout层、固定batchnorm的统计参数（验证时再开启）；把能离线完成的预处理（比如归一化）提前做好存入H5，避免在__getitem__里执行复杂变换。
• 换用更快的优化器：把SGD替换为AdamW或RAdam，这类优化器收敛速度更快，能减少达到目标精度所需的训练轮次。

四、优化数据预处理 #

• 分块存储预处理后的数据：如果内存放不下全量数据，把原H5拆成多个小H5文件（比如每个存10k张预处理后的图像），Dataset加载时并行读取多个小文件，提升读取效率。
• 用内置加速变换：优先用torchvision.transforms里的内置方法，这些方法基于C实现，比自定义Python函数快很多；必须自定义预处理时，用numpy或PyTorch张量操作替代纯Python循环。

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