从第一篇笔记写到现在,PyTorch 基础语法、张量操作、自动求导这些内容基本都过了一遍,但从"能跑通教程代码"到"能独立跑一个自己的训练任务",中间其实还隔着一段很长的路。这篇笔记 11 我打算换个思路,不再按部就班讲 API,而是把我自己在完整训练一个图像分类模型时踩过的坑、反复调整过的细节一次性梳理出来。内容围绕自定义 Dataset、DataLoader 参数调优、训练循环里的隐性细节、过拟合处理和多卡训练这几个方向展开,偏实战,适合已经掌握 PyTorch 基础、正在尝试跑自己数据的读者参考。
1. 从张量到手写 Dataset:我为什么建议你尽早摆脱内置数据集
1.1 内置数据集很好,但真实项目没有现成的
MNIST、CIFAR-10、ImageNet 这些内置数据集对学习来说非常友好,几行代码就能加载出来,但它们最大的问题在于——太干净了。文件名有序、标签正确、图片尺寸统一、没有损坏文件,所有脏活累活都被框架处理完了。一旦你换了真实场景的数据,比如自己从文件夹读图、处理 CSV 标注文件、过滤坏图,很多人会突然发现自己连数据加载都不会写了。
这不是危言耸听。我见过不少同学在 Kaggle 比赛里直接套用内置数据集的加载方式,然后被各种报错折磨:路径不对、标签错位、图片通道数不一致、显存爆掉还不知道是 DataLoader 一次拿了太多数据。所以这篇笔记第一个想强调的观点就是:尽早学会手写 Dataset,这是从学习走向实战的第一道分水岭。
1.2 Dataset 与 DataLoader 的完整实现:图片分类为例
PyTorch 的 Dataset 类核心只需要实现三个方法:__init__、__len__ 和 __getitem__。看起来简单,但实际写起来有不少细节决定着你后面能否顺利训练。
我以最常见的图像分类场景为例:图片存放在 train/images/ 下,标签存放在 train/label.csv 中,内容大概是 filename,label 两列。一个基础但完整的 Dataset 写法如下:
python复制import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
class ImageClassificationDataset(Dataset):
def __init__(self, img_dir, label_file, transform=None):
self.img_dir = img_dir
self.df = pd.read_csv(label_file)
# 建立类别名称到索引的映射
self.class_to_idx = {cls: idx for idx, cls in enumerate(sorted(self.df['label'].unique()))}
self.idx_to_class = {idx: cls for cls, idx in self.class_to_idx.items()}
self.transform = transform
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
row = self.df.iloc[idx]
img_path = os.path.join(self.img_dir, row['filename'])
# 用 try-except 捕获坏图,避免训练中断
try:
image = Image.open(img_path).convert('RGB')
except Exception:
# 坏图就随机拿一张正常的顶替
return self.__getitem__((idx + 1) % len(self.df))
label = self.class_to_idx[row['label']]
if self.transform:
image = self.transform(image)
return image, label
这个写法里有几个地方值得说明。第一,convert('RGB') 是必须的,因为真实数据集中经常混入灰度图、RGBA 图,不统一转 RGB 的话,后面 batch 的 shape 对不上,报错特别隐蔽。第二,坏图处理我用了一个简单粗暴的方案——跳过当前图片,取下一张。真实项目中更稳妥的做法是先在预处理阶段把所有坏图剔除掉,但训练过程中偶尔还是会遇到读取超时或文件被占用的情况,加一个 fallback 逻辑能避免整个训练任务直接挂掉。
1.3 一个容易被忽略的坑:返回的样本格式与 collate_fn 的关系
很多人第一次写自定义 Dataset 时,会在 __getitem__ 里返回一个字典:
python复制return {"image": image, "label": label, "filename": row['filename']}
这样做的优点是便于调试和扩展,但它有一个直接影响——DataLoader 默认的 collate_fn 会将字典中的每个
