PyTorch加载自己的数据集、搭一个LeNet5网络跑训练这活儿看着挺基础但真正动手的时候一堆细节能把人卡到怀疑人生图片尺寸对不上、通道数报错、全连接层维度不匹配、训练loss死活不降……我在带几个新人做图像分类项目时发现光一个“加载自己的数据”就能卡住大半天。这篇文章就把整个链路完整串一遍数据集目录怎么组织、什么时候写自定义Dataset、DataLoader和预处理怎么配、LeNet5在PyTorch里面怎么搭、训练循环怎么写以及最后那些容易踩的坑和排查技巧。适合刚学完深度学习基础、准备拿自己图片跑分类任务的人哪怕你现在手里只有几十张图也能跟着把流程跑通。1. 为什么选LeNet5 PyTorch这个组合1.1 LeNet5到底能干什么LeNet5是经典的手写数字识别网络但它的价值远不止识别数字。它是最早把卷积、池化、全连接组合起来的网络结构之一参数规模小、计算量低非常适合作为入门卷积神经网络的“活标本”。网络结构不复杂两个卷积层加三个全连接中间穿插池化层。输入是一张单通道的灰度图尺寸通常是32x32经过一系列卷积池化后把特征图展平最后输出10个类别的分数。正是因为它足够小用CPU跑都能很快看到训练结果所以特别适合用来验证“我的数据集加载对不对”“训练流程是否完整”。我经常建议新人在跑大模型之前先用这种小网络把数据和流程理顺不然数据集有问题都分不清是模型的问题还是数据的问题。1.2 PyTorch在这个场景下的优势选PyTorch不是因为它比别的框架“高级”而是因为它在这个场景下确实顺手。第一动态图机制让调试变得很直接。你在forward里面打印一行输出尺寸立刻就能定位到维度问题不用像静态图那样绕圈子。第二torchvision和torch.utils.data这套数据工具链很成熟ImageFolder、Dataset、DataLoader开箱即用省去自己造轮子的时间。第三代码写法和LeNet5论文里的结构能一一对应卷积、池化、全连接层堆叠起来非常直观对学习阶段的人来说很友好。还有一个实际原因PyTorch的社区资料多你在训练中遇到任何报错基本都能搜到别人是怎么解决的。这对于新手来说比任何理论讲解都管用。2. 数据准备先把自己的数据集理顺2.1 推荐用文件夹目录结构组织分类数据对于图片分类任务数据组织有个“最省力方案”按类别建文件夹文件夹名就是标签名。这就是torchvision里ImageFolder期望的格式。data/ ├── train/ │ ├── cat/ │ │ ├── cat_001.jpg │ │ ├── cat_002.jpg │ │ └── ... │ └── dog/ │ ├── dog_001.jpg │ └── ... └── val/ ├── cat/ └── dog/这种结构的优点有两个。一是可以不用手写Dataset类直接datasets.ImageFolder就能加载少写代码少出错。二是标签顺序按文件夹名字的字母顺序自动生成cat对应0、dog对应1逻辑透明不容易搞混。如果你只是做练习不追求严格的训练集和验证集划分那至少也要准备两批数据一批用来训练一批用来评估。千万不要把训练数据和评估数据混在一起否则你看不到模型对新图片的真实泛化效果。2.2 什么时候必须自己写Dataset类ImageFolder方便但很多真实场景没法用文件夹结构表达。比如你的样本信息存在CSV文件里每行是“图片路径, 标签”或者你做的不是单标签分类而是多标签任务再或者你需要对文件名做复杂解析。这些时候就要继承torch.utils.data.Dataset自定义数据集。我自己常用的写法是import os from PIL import Image from torch.utils.data import Dataset class MyDataset(Dataset): def __init__(self, root_dir, transformNone): self.samples [] self.transform transform classes sorted(os.listdir(root_dir)) self.class_to_idx {cls_name: idx for idx, cls_name in enumerate(classes)} self.classes classes for class_name in classes: class_dir os.path.join(root_dir, class_name) for img_name in os.listdir(class_dir): img_path os.path.join(class_dir, img_name) self.samples.append((img_path, self.class_to_idx[class_name])) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] image Image.open(img_path).convert(L) if self.transform: image self.transform(image) return image, label这段代码里有两个细节需要注意。第一convert(L)把图片转成灰度单通道配合LeNet5原始的1个输入通道。如果你的图片本身就是灰度图那加这一句话没有副作用如果图片是彩色图不转的话就会变成3通道而模型第一个卷积层输入是1通道两者对不上运行时会直接报错。第二__getitem__里返回的是PIL图片和标签真正转Tensor的操作交给transform做。这么做的好处是灵活你想调整预处理方案改transform就行不用动数据读取代码。2.3 数据集划分和文件检查动手训练之前先做几件简单但重要的事。第一统计一下每个类别的图片数量。如果有类别数量差异极大比如一类有500张、另一类只有50张训练出来的模型大概率偏向多的那一类。第二随机抽几十张图肉眼看一下图片内容是否和标签匹配目录里混入损坏图片或者文字截图的情况我见过不止一次。第三确认所有图片都能正常打开最好写个小脚本遍历一遍。from PIL import Image import os bad_images [] for root, _, files in os.walk(data/train): for f in files: path os.path.join(root, f) try: img Image.open(path) img.load() except Exception: bad_images.append(path) print(f损坏图片数量: {len(bad_images)})这一步能帮你省下训练过程中突然报Image file is truncated或者解压失败的麻烦。3. 预处理、Dataset与DataLoader的完整实现3.1 预处理流程不是只有ToTensor从图片文件到模型输入中间要经过三个关键转换尺寸统一、数据类型转换、归一化。很多人只写了ToTensor()就开跑结果模型要么收敛慢要么直接不收敛。尺寸统一这块LeNet5的原始输入是32x32所以最好统一到32x32。可以用Resize直接缩放到32x32也可以用Pad在四周补零把任意尺寸的图变成32x32。Resize简单直接Pad会保留图片原始比例和内容但会增加黑色边框区域。我个人更建议先短边Resize到固定尺寸、再做中心裁剪但这只是习惯问题对LeNet5这种网络的最终影响不会太大。ToTensor()做的事值得多说一句它把PIL图片或者numpy数组转成Tensor同时把像素值从0到255归一化到0到1。这时候数据的数值范围已经适合训练了。至于Normalize(mean, std)它对每个通道做标准化让数据分布更接近标准正态分布有助于稳定训练。如果你用的是灰度图均值填0.5、标准差填0.5就够了如果你处理的是彩色图就要填三个通道的值。也可以用数据集的真实统计值来算但这里不必过度纠结。from torchvision import transforms transform transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ])有一点要注意这三个步骤的顺序不能乱。Resize必须在ToTensor之前因为ToTensor之后数据就不是PIL图片了而Normalize必须在ToTensor之后因为它操作的是Tensor数组。如果你希望模型泛化能力更强一点还可以在训练集上加一些数据增强比如随机旋转、随机水平翻转、颜色抖动之类。但LeNet5网络容量有限如果数据集本身就很小过度的数据增强反而不利于收敛适度即可。3.2 DataLoader参数设置每个参数都有讲究数据准备完毕接下来用DataLoader把数据集包一层。from torch.utils.data import DataLoader train_dataset MyDataset(data/train, transformtransform) val_dataset MyDataset(data/val, transformtransform) train_loader DataLoader( train_dataset, batch_size64, shuffleTrue, num_workers2, pin_memoryTrue ) val_loader DataLoader( val_dataset, batch_size64, shuffleFalse, num_workers2, pin_memoryTrue )batch_size决定了每次迭代送入模型的样本数量。调大batch会让训练更稳定但对显存要求更高调小batch会让梯度更新更频繁收敛过程会更震荡。LeNet5很小64或者128都可以我习惯从64起步。shuffle在训练集必须设置为True在验证集和测试集设置为False。这样每个epoch训练数据的顺序都会打乱避免模型学到样本顺序的规律验证集不需要打乱因为评估结果不应该受顺序影响。num_workers控制数据加载的进程数设为0表示只用主进程加载Windows下不报错但速度慢。调到2或者4能明显加快数据读取。不过要注意Windows系统上如果没把数据加载代码放在if __name__ __main__:保护块里num_workers0会反复启动子进程甚至直接崩溃。pin_memoryTrue可以加快GPU训练时的数据传输效率前提是你用CUDA训练。纯CPU环境没有明显收益但开了也不会有大问题。3.3 训练前先做一次数据体检我有个习惯正式训练前先手动取一个batch看看数据长什么样。images, labels next(iter(train_loader)) print(图片张量形状:, images.shape) print(标签张量形状:, labels.shape) print(标签范围:, labels.min().item(), labels.max().item())如果images.shape是[64, 1, 32, 32]说明格式对了如果是[64, 3, 32, 32]说明通道没处理好。标签范围也应该在0到类别数减1之间。这一步能提前暴露80%的维度和标签问题比训练到一半再报错强得多。4. LeNet5模型搭建从结构推演到PyTorch实现4.1 先把每一层输出尺寸算清楚LeNet5的结构可以拆成两部分特征提取部分和分类部分。以32x32的灰度输入为例逐层推导过程如下。第一个卷积层用6个5x5的卷积核没有填充输出尺寸是32-5128得到6个28x28的特征图。接着2x2池化尺寸减半到14x14。第二个卷积层用16个5x5卷积核输出10x10的特征图。再经过一次2x2池化得到16个5x5的特征图。最后再接一个5x5卷积卷积核数量是120输出就是120个1x1的数值。展平后进入全连接层先映射到84维再映射到类别数。我把完整尺寸变化整理成下面的表格网络层输出尺寸参数量公式输入1x32x32-Conv2d(1, 6, 5)6x28x28(1x5x51)x6156MaxPool2d(2)6x14x14-Conv2d(6, 16, 5)16x10x10(6x5x51)x162416MaxPool2d(2)16x5x5-Conv2d(16, 120, 5)120x1x1(16x5x51)x12048120Flatten120-Linear(120, 84)84120x848410164Linear(84, num_classes)num_classes84x1010850这一套算下来总参数量也就六万出头在今天的标准里算是很轻量的网络。很多人问如果输入图不是32x32怎么办其实只要你在预处理阶段统一Resize成32x32就不用操心尺寸问题。如果不改输入尺寸比如直接用28x28的原生MNIST那么最后一个卷积层会因为没有足够空间卷积而报错你需要重新计算全连接层的维度。与其这样绕来绕去不如老老实实统一输入尺寸。4.2 PyTorch代码实现按上面的结构写代码一层对应一行清晰直观。import torch import torch.nn as nn class LeNet5(nn.Module): def __init__(self, num_classes10): super(LeNet5, self).__init__() self.features nn.Sequential( nn.Conv2d(1, 6, kernel_size5, stride1), # 32x32 - 28x28 nn.ReLU(), nn.MaxPool2d(kernel_size2, stride2), # 28x28 - 14x14 nn.Conv2d(6, 16, kernel_size5, stride1), # 14x14 - 10x10 nn.ReLU(), nn.MaxPool2d(kernel_size2, stride2), # 10x10 - 5x5 nn.Conv2d(16, 120, kernel_size5, stride1), # 5x5 - 1x1 nn.ReLU() ) self.classifier nn.Sequential( nn.Linear(120, 84), nn.ReLU(), nn.Linear(84, num_classes) ) def forward(self, x): x self.features(x) x torch.flatten(x, 1) x self.classifier(x) return x和原版LeNet5相比这里做了两个小改动。第一原论文使用的是sigmoid激活函数现在普遍用ReLU收敛更快代码里可以保留ReLU。第二原论文最后用的是欧氏距离径向基函数做分类现在我们直接换成全连接层和Softmax交叉熵这是目前更主流的做法也方便和PyTorch的CrossEntropyLoss搭配。如果你的数据集是彩色图就把nn.Conv2d(1, 6, ...)改成nn.Conv2d(3, 6, ...)其余不变。4.3 模型实例化和参数统计定义好模型之后随手加两行代码确认结构没问题。model LeNet5(num_classes10) print(model) total_params sum(p.numel() for p in model.parameters()) print(f总参数量: {total_params})torch.flatten(x, 1)从第一维开始展平保持batch维度不变。假如输入是[64, 120, 1, 1]展平后变成[64, 120]正好灌进全连接层。还有个容易写错的地方如果在forward里用了x.view(x.size(0), -1)效果和flatten一样但view要求内存连续如果之前做过转置之类的操作可能会报错。flatten更稳妥。5. 训练主循环从损失函数到模型保存5.1 损失函数和优化器怎么选分类任务的标准组合就是CrossEntropyLoss加优化器。CrossEntropyLoss内部把Softmax和交叉熵合并了所以你不需要在模型输出层额外加Softmax。模型输出的原始分数直接送进损失函数即可。优化器最经典的选择是SGD加动量criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9)这里lr0.01和momentum0.9是很多LeNet参考实现里比较稳的起点。SGD加动量训练起来稳定、收敛性能好但需要手动调节学习率。如果嫌调学习率麻烦也可以用Adamoptimizer torch.optim.Adam(model.parameters(), lr0.001)Adam自带自适应学习率对于简单任务通常能更快见到效果但有时候最终精度不如调好的SGD。我的建议是第一次跑通流程用Adam因为踩坑少等到追求最好效果时再切回SGD仔细调学习率。5.2 完整训练循环代码一个标准的训练过程包括训练模式和验证模式两段。训练模式开启梯度计算和权重更新验证模式关闭梯度计算只做前向推理。import torch device torch.device(cuda if torch.cuda.is_available() else cpu) model LeNet5(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9) epochs 20 best_acc 0.0 for epoch in range(epochs): # 训练阶段 model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() train_loss running_loss / total train_acc correct / total # 验证阶段 model.eval() val_loss 0.0 val_correct 0 val_total 0 with torch.no_grad(): for images, labels in val_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) val_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) val_total labels.size(0) val_correct (predicted labels).sum().item() val_loss val_loss / val_total val_acc val_correct / val_total print(fEpoch {epoch1:3d}/{epochs} | fTrain Loss: {train_loss:.4f} Acc: {train_acc:.4f} | fVal Loss: {val_loss:.4f} Acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f 最佳模型已保存验证准确率: {best_acc:.4f})这段代码里有一个关键操作需要说明optimizer.zero_grad()必须在计算梯度之前调用。PyTorch的梯度是累加式的如果不每次清空梯度会在多个batch之间累加那训练行为就和预期完全不一样了。with torch.no_grad()是验证阶段的标准写法它告诉PyTorch不要跟踪梯度既省内存又加快速度。model.eval()则是把模型切换到推理模式影响的主要是BatchNorm和Dropout这类层。LeNet5里面暂时没有这些层但养成写eval()和train()切换的习惯很重要以后迁移到复杂网络才不会漏。5.3 模型保存、加载与单张预测保存模型时我建议保存state_dict而不是整个模型对象。state_dict是权重字典体积小、可移植性强加载时重新构建模型结构再灌入权重。上面代码已经把最佳模型的权重存到best_model.pth下面演示如何加载。model LeNet5(num_classes10) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.to(device) model.eval()单张图片预测时前处理步骤必须和训练时完全一致灰度转换、Resize到32x32、转Tensor、走同样的Normalize。from PIL import Image import torchvision.transforms as transforms def predict(img_path): transform transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) image Image.open(img_path).convert(L) tensor transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): output model(tensor) pred output.argmax(dim1).item() return pred注意unsqueeze(0)这一步。训练时数据是[batch, channel, height, width]单张图没有batch维度所以要手动加一个维度否则模型会报错。6. 训练踩坑实录与调参经验6.1 高频报错与排查方案这部分整理的是我在跑LeNet5和数据加载时最常遇到的错误。第一个是维度不匹配报错信息类似Expected input batch_size to match target batch_size或者mat1 and mat2 shapes cannot be multiplied。这个大概率是全连接层输入尺寸算错了。解决办法很简单在模型forward里临时打印每一层输出的形状肉眼就能看出问题。第二个是图片通道数对不上报错信息是Given groups1, weight of size [6, 1, 5, 5], expected input to have 1 channel, but got 3 channels。这就是彩色图没转灰度。解决办法是在Dataset的__getitem__里加convert(L)或者把第一个卷积层输入通道改成3。第三个是DataLoader子进程崩溃用num_workers0时在Windows上报错或程序重启。解决方法是把数据加载和训练逻辑包在if __name__ __main__:中或者临时把num_workers设为0先跑通。第四个是图片损坏报错OSError: image file is truncated。在Image.open时加上参数能缓解from PIL import ImageFile ImageFile.LOAD_TRUNCATED_IMAGES True但最好还是提前检查数据把损坏图片清理掉治标不治本只会让后续干扰训练。第五个是loss不收敛。如果训练集准确率上不去先怀疑学习率太大或太小然后检查数据归一化和标签对不对。如果训练集准确率很高但验证集准确率低那是过拟合需要加数据增强、加正则化或者减小模型容量。我把这些整理成下面的速查表现象可能原因排查方向维度报错全连接层尺寸不符合特征图尺寸打印每层输出形状逐一核对通道数报错彩色图未转灰度或in_channels写错检查Dataset的convert(L)和Conv2d第一层Windows下DataLoader崩溃num_workers0且未加main保护包if __name__ __main__或临时设0图片读取报错图片文件损坏或格式异常扫描数据集删除损坏文件训练loss震荡不降学习率过大适当降低学习率观察是否稳定验证集准确率低过拟合或验证集分布偏差加数据增强、早停、检查验证集6.2 小数据集怎么训练LeNet5LeNet5毕竟是个小网络如果你的训练集只有几百张图有几条经验值得留意。第一别让数据增强变成破坏。比如随机水平翻转对猫狗这类任务没问题但对手写数字这种方向敏感的图片就很糟糕。增强手段一定要跟数据特点匹配。第二用验证集做“温度计”。每训练一个epoch看一眼验证准确率如果验证准确率连续几个epoch不涨甚至下降就考虑早停别傻傻跑完所有epoch。第三适当降低batch_size比如从64降到32让每个batch的梯度估计更频繁地更新在小数据上反而更稳。还有一种情况经常遇到就是训练集很小、但类别很多。这时候模型很容易把验证集里没见过的样本分错。最直接的办法是尽量收集更多数据其次是用预训练模型做迁移学习但这就不是LeNet5的适用范围了。如果必须用LeNet5可以试着减少后面的全连接层宽度降低过拟合风险。说实话LeNet5天然就不是为大数据准备的小规模分类、教学演示、流程验证才是它最舒服的舞台。6.3 提升训练效果的个人心得从我自己跑过的多次实验看有五个点对最终效果影响最大。第一归一化参数别拍脑袋。虽然用0.5和0.5能跑通但如果你把Normalize的均值标准差改成数据集的真实统计值训练速度和最终精度通常会有小幅提升。第二学习率值得花时间调。LeNet5配合SGD时0.01是个普遍有效的起点但如果你发现loss下降非常慢试着用0.05或者0.001对比一下。第三训练过程中把loss曲线画出来比只看准确率能发现更多问题。第四保存模型时连优化器状态一起保存这样恢复训练时能无缝接着跑。第五也是最容易被忽略的数据和标签顺序要稳定。每次运行代码之前把数据集目录字段和标签映射打印出来确认一下防止因为文件夹改名导致标签错位。6.4 用TensorBoard记录训练曲线如果你不想只看打印的日志PyTorch官方工具TensorBoard可以帮你把loss和accuracy画成曲线。pip install tensorboard在训练代码里引入SummaryWriterfrom torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/lenet5_experiment) # 每个epoch结束后记录 writer.add_scalar(Loss/train, train_loss, epoch) writer.add_scalar(Accuracy/train, train_acc, epoch) writer.add_scalar(Loss/val, val_loss, epoch) writer.add_scalar(Accuracy/val, val_acc, epoch)训练结束后在终端运行tensorboard --logdirruns/lenet5_experiment浏览器打开终端提示的地址就能看到曲线。这个习惯一旦养成以后跑任何模型都会很方便因为你不需要等到训练全部结束才判断模型好坏训练中途就能看到趋势。7. 写在最后的几句实在话LeNet5这个网络放在今天性能早就不是最先进的但它作为自己动手写数据加载和训练流程的第一个模型地位几乎没有替代品。我自己每次接触新的数据集都习惯先拿这类小网络跑一个baseline验证数据链路通了、模型训练能正常收敛再考虑上更复杂的结构。最后再分享一个小技巧如果你的数据是灰度图、尺寸又都比较小用LeNet5做验证实验时CPU完全够用不需要一开始就上GPU。这样反而能更清晰地理解每一行代码在做什么而不是被显存、批大小这些工程问题分散注意力。等你在小数据上把整个流程跑熟再迁移到大数据、大模型你会发现当初这些看似琐碎的细节全都是后来调试模型的底气。