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, transform=None): 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', transform=transform) val_dataset = MyDataset('data/val', transform=transform) train_loader = DataLoader( train_dataset, batch_size=64, shuffle=True, num_workers=2, pin_memory=True ) val_loader = DataLoader( val_dataset, batch_size=64, shuffle=False, num_workers=2, pin_memory=True )batch_size决定了每次迭代送入模型的样本数量。调大batch会让训练更稳定，但对显存要求更高；调小batch会让梯度更新更频繁，收敛过程会更震荡。LeNet5很小，64或者128都可以，我习惯从64起步。
shuffle在训练集必须设置为True，在验证集和测试集设置为False。这样每个epoch训练数据的顺序都会打乱，避免模型学到样本顺序的规律；验证集不需要打乱，因为评估结果不应该受顺序影响。
num_workers控制数据加载的进程数，设为0表示只用主进程加载，Windows下不报错但速度慢。调到2或者4能明显加快数据读取。不过要注意，Windows系统上如果没把数据加载代码放在if __name__ == '__main__':保护块里，num_workers>0会反复启动子进程，甚至直接崩溃。
pin_memory=True可以加快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-5+1=28，得到6个28x28的特征图。接着2x2池化，尺寸减半到14x14。第二个卷积层用16个5x5卷积核，输出10x10的特征图。再经过一次2x2池化，得到16个5x5的特征图。最后再接一个5x5卷积，卷积核数量是120，输出就是120个1x1的数值。展平后进入全连接层，先映射到84维，再映射到类别数。
我把完整尺寸变化整理成下面的表格：
| 网络层 | 输出尺寸 | 参数量公式 |
|---|---|---|
| 输入 | 1x32x32 | - |
| Conv2d(1, 6, 5) | 6x28x28 | (1x5x5+1)x6=156 |
| MaxPool2d(2) | 6x14x14 | - |
| Conv2d(6, 16, 5) | 16x10x10 | (6x5x5+1)x16=2416 |
| MaxPool2d(2) | 16x5x5 | - |
| Conv2d(16, 120, 5) | 120x1x1 | (16x5x5+1)x120=48120 |
| Flatten | 120 | - |
| Linear(120, 84) | 84 | 120x84+84=10164 |
| Linear(84, num_classes) | num_classes | 84x10+10=850 |
这一套算下来总参数量也就六万出头，在今天的标准里算是很轻量的网络。
很多人问，如果输入图不是32x32怎么办？其实只要你在预处理阶段统一Resize成32x32，就不用操心尺寸问题。如果不改输入尺寸，比如直接用28x28的原生MNIST，那么最后一个卷积层会因为没有足够空间卷积而报错，你需要重新计算全连接层的维度。与其这样绕来绕去，不如老老实实统一输入尺寸。
4.2 PyTorch代码实现
按上面的结构写代码，一层对应一行，清晰直观。
import torch import torch.nn as nn class LeNet5(nn.Module): def __init__(self, num_classes=10): super(LeNet5, self).__init__() self.features = nn.Sequential( nn.Conv2d(1, 6, kernel_size=5, stride=1), # 32x32 -> 28x28 nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), # 28x28 -> 14x14 nn.Conv2d(6, 16, kernel_size=5, stride=1), # 14x14 -> 10x10 nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), # 10x10 -> 5x5 nn.Conv2d(16, 120, kernel_size=5, stride=1), # 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_classes=10) 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(), lr=0.01, momentum=0.9)这里lr=0.01和momentum=0.9是很多LeNet参考实现里比较稳的起点。SGD加动量训练起来稳定、收敛性能好，但需要手动调节学习率。
如果嫌调学习率麻烦，也可以用Adam：
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)Adam自带自适应学习率，对于简单任务通常能更快见到效果，但有时候最终精度不如调好的SGD。我的建议是，第一次跑通流程用Adam，因为踩坑少；等到追求最好效果时再切回SGD仔细调学习率。
5.2 完整训练循环代码
一个标准的训练过程包括训练模式和验证模式两段。训练模式开启梯度计算和权重更新，验证模式关闭梯度计算，只做前向推理。
import torch device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = LeNet5(num_classes=10).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.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(f"Epoch {epoch+1:3d}/{epochs} | " f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | " f"Val 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_classes=10) model.load_state_dict(torch.load('best_model.pth', map_location=device)) 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(dim=1).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 groups=1, weight of size [6, 1, 5, 5], expected input to have 1 channel, but got 3 channels。这就是彩色图没转灰度。解决办法是在Dataset的__getitem__里加convert('L')，或者把第一个卷积层输入通道改成3。
第三个是DataLoader子进程崩溃，用num_workers>0时在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_workers>0且未加main保护 | 包if __name__ == '__main__'或临时设0 |
| 图片读取报错 | 图片文件损坏或格式异常 | 扫描数据集，删除损坏文件 |
| 训练loss震荡不降 | 学习率过大 | 适当降低学习率，观察是否稳定 |
| 验证集准确率低 | 过拟合或验证集分布偏差 | 加数据增强、早停、检查验证集 |
6.2 小数据集怎么训练LeNet5
LeNet5毕竟是个小网络，如果你的训练集只有几百张图，有几条经验值得留意。
第一，别让数据增强变成破坏。比如随机水平翻转，对猫狗这类任务没问题，但对手写数字这种方向敏感的图片就很糟糕。增强手段一定要跟数据特点匹配。第二，用验证集做“温度计”。每训练一个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在训练代码里引入SummaryWriter：
from 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 --logdir=runs/lenet5_experiment浏览器打开终端提示的地址就能看到曲线。这个习惯一旦养成，以后跑任何模型都会很方便，因为你不需要等到训练全部结束才判断模型好坏，训练中途就能看到趋势。
7. 写在最后的几句实在话
LeNet5这个网络放在今天，性能早就不是最先进的，但它作为自己动手写数据加载和训练流程的第一个模型，地位几乎没有替代品。我自己每次接触新的数据集，都习惯先拿这类小网络跑一个baseline，验证数据链路通了、模型训练能正常收敛，再考虑上更复杂的结构。
最后再分享一个小技巧：如果你的数据是灰度图、尺寸又都比较小，用LeNet5做验证实验时CPU完全够用，不需要一开始就上GPU。这样反而能更清晰地理解每一行代码在做什么，而不是被显存、批大小这些工程问题分散注意力。等你在小数据上把整个流程跑熟，再迁移到大数据、大模型，你会发现当初这些看似琐碎的细节，全都是后来调试模型的底气。
