资讯动态

PyTorch彩色图片分类实战:CIFAR-10数据管道与CNN训练调参

发布时间:2026/10/22 7:51:20 来源:尧图企业网站定制
老实说看到“第T2周彩色图片分类”这个任务的时候我第一反应是这不就是把上周的灰度图片换个通道数而已吗等真的开始写数据加载和训练脚本才发现彩色图片分类的坑一半在模型定义之前就已经埋好了。这篇文章是给正在做类似周任务的读者准备的尤其适合那些已经能跑通基础训练流程、但还没系统梳理过彩色图像分类关键环节的人。我会用经典的小型彩色图片数据集 CIFAR-10 来做演示把数据加载、归一化、数据增强、卷积网络设计、训练验证、调参踩坑这几个环节完整过一遍。如果你想换成自己的彩色图片数据集大部分逻辑保持不变只需要改一下数据集入口和类别数量。1. 为什么“彩色图片分类”比灰度任务麻烦这么多1.1 输入从1个通道变成3个通道模型和内存一起变灰度图在送入神经网络时张量形状一般是(batch, 1, H, W)模型的第一层卷积输入通道数是 1。彩色图则不同一张 RGB 图的形状是(height, width, 3)进入 PyTorch 后要变成(batch, 3, H, W)。这不只是把第一个卷积层的in_channels改成 3 那么简单。3 个通道意味着模型需要在同一个位置上同时看到三个颜色平面的联合信息。换句话说模型要学的不是“三张灰度图分别是什么”而是“三张图叠在一起时这个颜色分布对应哪个类别”。如果只是把输入通道改成 3其他结构不变实际计算量会从输入层开始增加一部分但对整体参数量影响不大真正变化的是数据管道的复杂程度和归一化的方式。我在第一周跑灰度分类时模型里的输入张量是(1, 28, 28)这类结构从未考虑过通道顺序。到了彩色图片任务第一次写transforms.Normalize的时候我差点把三通道均值错写成三个相同的数。这也是新手最容易踩的坑之一彩色图的三通道有各自的均值和标准差不能像灰度图那样当成一个标量来处理。1.2 数据集与标签经典十类别彩色小图这次周任务我选的是 CIFAR-106 万张 32×32 的彩色图片10 个类别每个类别 6000 张其中 5 万张训练、1 万张测试。每张图都很小用 CPU 也能勉强训练但真要跑出好看的效果还是建议开 GPU。类别列表如下标签索引类别名标签索引类别名0airplane5dog1automobile6frog2bird7horse3cat8ship4deer9truck它有很多预先做好的 loader我直接借助torchvision来加载省去下载后手工整理目录的麻烦。对于自建数据集建议先看一眼图片读取出来的通道顺序到底是 RGB 还是 BGR。很多图像库默认加载顺序不一样这一步错了后面所有颜色相关的结果都会莫名其妙。2. 数据管道加载、归一化、增强的先后顺序不能乱2.1 用 torchvision 一口气搞定加载与预处理我最开始写数据加载时习惯直接用PIL循环读文件夹代码又长又容易出错。这周的彩色图片分类任务我换成了torchvision.datasets加transforms.Compose的方案几行代码就能把读取、随机裁剪、随机翻转、转张量、归一化串起来。import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) ]) transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) ]) train_dataset datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) test_dataset datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers2, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_size128, shuffleFalse, num_workers2, pin_memoryTrue)这里的顺序是有讲究的RandomCrop和RandomHorizontalFlip必须放在ToTensor之前因为它们操作的对象是 PIL 图片Normalize必须放在ToTensor之后因为它要对浮点型张量按通道做计算。2.2 均值方差计算与归一化的真实作用CIFAR-10 的归一化参数不是随便写的而是对整个训练集每一张图的三通道像素分别统计出来的均值和方法。以第一个通道为例均值 0.4914 意味着如果把所有训练图的 R 通道像素取平均再除以 255结果大约是 0.4914。归一化公式就是(x - mean) / std。如果跳过这一步模型输入像素值都在 0 到 1 之间分布集中训练前期梯度的量级会偏大学习率稍微给高一点就会出现 loss 反复震荡。做了归一化之后每一维输入大致呈均值接近 0、方差接近 1 的分布网络收敛会明显平稳得多。如果你手里的彩色图片不是 CIFAR-10我建议花几分钟自己统计一次均值和标准差而不是直接抄别人的参数。做法也很简单把所有训练图片缩放到统一尺寸后逐图片累加每个通道的像素和像素平方最后算均值和方差。数据量大时可以用采样一部分图片来近似算通常 2000 到 5000 张就足够稳定了。2.3 增强的取舍RandomCrop 和 RandomHorizontalFlip我这次用了两种数据增强RandomCrop(32, padding4)和RandomHorizontalFlip()。前者会先把图片四周各补 4 个像素再随机裁剪回 32×32相当于让模型每次看到同一张图的不同局部区域对缓解过拟合很有帮助后者把图片随机水平翻转相当于把数据量近似翻倍。但数据增强不是越多越好。如果你的图片分类本身对方向敏感比如要区分“朝左”和“朝右”的物体水平翻转就会把语义完全弄反。CIFAR-10 里的类别对翻转基本不敏感所以水平翻转在这里是安全的。我见过很多人把RandomCrop写成了RandomResizedCrop。那样做随机缩放的尺度变化很大会对小物体的语义信息产生破坏在 32×32 这种小图上尤其明显。除非你确定原图分辨率高、物体占比较大否则不要盲目把增强堆得很复杂。3. 自己搭一个够用的小型卷积网络3.1 结构设计小卷积核叠加和池化下采样彩色图片分类最直接的模型是卷积神经网络。很多人一开始会想用大卷积核比如 7×7觉得感受野越大越好。但在 32×32 的输入图上大卷积核会把空间尺寸快速压小而且参数量不低。我这次用的小网络完全由 3×3 卷积组成基本结构是“两个 3×3 卷积 批归一化 ReLU 最大池化”为一组重复三组再接全连接分类器。为什么用小卷积核叠加因为两个 3×3 卷积堆叠起来感受野相当于一个 5×5 卷积但参数量更少中间还多了一次非线性变换表达能力更强这是实际项目中很常用的设计思路。输入图片 32×32经过三次最大池化后空间尺寸依次变成 16×16、8×8、4×4。最后一层卷积输出 64 个通道所以展平后的特征向量长度是64×4×41024。3.2 关键层不是越多越好BatchNorm、Dropout 的定位网络里我加了BatchNorm2d和Dropout这两个层看起来不增加特征提取能力却对训练过程有决定性影响。BatchNorm2d的作用不是让网络更深而是让每一层的激活值分布尽量稳定。它一方面允许你用稍大一点的学习率另一方面也能减轻梯度在深层传播时的异常变化。要注意的是BatchNorm 在训练和预测时的行为不一样训练时使用当前批次的统计量验证时使用累计的全局统计量。这也是后面必须写model.train()和model.eval()的原因之一。Dropout我用在全连接层附近比例分别是 0.25 和 0.5。它的作用是随机丢掉一部分神经元迫使网络不把希望全部寄托在某几个节点上是典型的抗过拟合手段。卷积层后面通常不用 Dropout因为 BatchNorm 已经能承担一部分正则化的作用。3.3 完整模型代码与参数量下面是这周任务里的完整模型代码输入是(batch, 3, 32, 32)输出是(batch, 10)import torch import torch.nn as nn class ColorCNN(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, 3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.Conv2d(16, 16, 3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.Conv2d(32, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.Conv2d(64, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Dropout(0.25), nn.Linear(64 * 4 * 4, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) x self.classifier(x) return x写完模型后我一般会先用一个假张量检查形状有没有问题model ColorCNN() fake_input torch.zeros(2, 3, 32, 32) output model(fake_input) print(output.shape) # 期望输出 torch.Size([2, 10])这个模型的参数量大概在 34 万左右。它不算大放在 CPU 上也能训练但比上周那种几个全连接层的模型明显能学到更丰富的颜色和纹理特征。4. 训练与验证从 loss 到 test accuracy 的完整闭环4.1 训练循环的标准写法训练函数我习惯写成下面这样的结构它负责一个 epoch 的完整遍历并返回训练集上的平均 loss 和准确率。def train_one_epoch(model, data_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in data_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) correct (outputs.argmax(1) labels).sum().item() total labels.size(0) epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc验证函数写法类似但有一个明显区别验证阶段不需要反向传播需要把梯度全部关闭并且把模型切成评估模式。def evaluate(model, data_loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in data_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) running_loss loss.item() * images.size(0) correct (outputs.argmax(1) labels).sum().item() total labels.size(0) epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc完整训练时我会在每个 epoch 结束后同时打印训练和测试结果方便实时观察两者的差距。如果训练准确率快速上涨而测试准确率停滞不前说明过拟合开始出现了。4.2 优化器与学习率的选择分类问题的损失函数用的是nn.CrossEntropyLoss()它在内部已经把LogSoftmax和NLLLoss结合在一起了所以模型的最后一层不需要手动加 softmax直接输出 10 个原始分数即可。只有在做推理、想要概率分布的时候才需要对输出做softmax。优化器我这次用 Adam学习率0.001。对小型卷积网络来说Adam 的开箱体验不错基本不需要太精细的预热策略。如果你更喜欢传统一点的做法也可以换成带动量的 SGD例如SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4)。两种优化器的区别在经验上比较明显Adam 前期收敛快但接近最优解时可能不稳定SGD 前期慢但调好学习率后最终精度往往更稳。做周任务的话我建议直接用 Adam省心。如果之后追求更好的测试准确率再切到 SGD 配合学习率衰减策略。4.3 验证阶段最容易漏掉的 eval我见过最多的训练事故之一就是只在训练函数里写了model.train()到了验证环节忘了写model.eval()。这个错误在纯全连接层网络里可能看不出来但只要模型里有 BatchNorm 或 Dropout结果就会很奇怪。验证时不写model.eval()Dropout 会继续随机丢弃神经元导致每次验证的结果都不一样BatchNorm 则会继续用当前 batch 的统计量而不是训练阶段累积的全局统计量预测的稳定性会下降。所以训练函数和验证函数里的模式切换必须成对出现这个习惯越早养成越好。5. 踩坑排查链路跑不出分数时先查这几处5.1 先让模型过拟合一个小批量验证网络没写错如果训练了好几轮 loss 都不降别急着改结构先做一个最基础的检查拿一个 batch 的数据反复在同一个 mini-batch 上训练几十步。如果网络本身没有明显 bugloss 应该能快速降到一个很小的值训练集准确率可能直接变成 100%。model ColorCNN() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.001) images, labels next(iter(train_loader)) for step in range(50): optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() if (step 1) % 10 0: acc (outputs.argmax(1) labels).float().mean().item() print(fstep {step 1}: loss{loss.item():.4f} acc{acc:.4f})如果几十步下来 loss 还在原地打转那就不是数据或超参的问题而是模型定义、输入预处理或 loss 计算方式本身写得不对。这种“故意过拟合一个小批量”的调试思路能帮你把网络结构问题和训练策略问题分离开来。5.2 归一化参数和通道顺序错位彩色图片分类的另一个高频坑是通道顺序。比如用 OpenCV 读图默认是 BGR 顺序而训练时使用的归一化均值是 RGB 顺序那么红色通道上减的其实是蓝色通道的均值颜色信息和语义标签就会被错误地绑定。我建议在数据加载后统一打印一次输入张量的形状并检查第一个通道的最小值和最大值排除通道顺序问题。如果transforms.ToTensor()之后得到的是(3, H, W)且三通道数值分布合理那通道层面基本没问题。有时候Normalize用错了维数也会出问题。灰度图可以写Normalize((0.5,), (0.5,))但彩色图必须写三个通道对应的均值和标准差。如果填成了Normalize(0.5, 0.5)PyTorch 会直接报维度错误不过这一般比静默出错要好处理一些。5.3 训练/验证模式切换与 BatchNorm 统计量除了model.eval()没写还有一个比较隐蔽的问题是验证集顺序。即使你写了model.eval()如果每次验证都随机打乱数据两个 epoch 之间的验证结果仍然会有波动不利于判断模型是否真的变好了。因此验证集的 DataLoader 我通常固定shuffleFalse。如果发现训练集 loss 已经很低但测试集准确率还上不去甚至训练准确率和测试准确率差距越来越大大概率是过拟合。此时优先做的事是增加数据增强强度比如再引入随机旋转或者加一点点随机擦除。其次可以增加 Dropout 比例。不要一开始就盲目加大模型容量模型越大在固定数据集上越容易过拟合。5.4 数据加载速度与进程问题很多人训练到一半发现每个 epoch 之间卡顿严重以为模型太复杂其实瓶颈在数据加载。DataLoader里的num_workers可以设置并行加载进程数pin_memoryTrue可以加速数据从内存搬到 GPU 的过程。不过num_workers并不是越大越好设置过大容易直接吃掉大量内存。在 Windows 环境里num_workers大于 0 时偶尔会触发多进程相关的报错。如果你用的是 Windows 且不知道怎么处理可以把num_workers改成 0速度慢一些但至少稳定。如果你的 CPU 核数够多建议num_workers从 2 到 4 开始试。6. 结果怎么看准确率之外的单类表现6.1 训练曲线与 test accuracy 的合理区间这个小型卷积网络在 CIFAR-10 上如果只做基础训练不加增强20 到 30 个 epoch 之后测试准确率通常在 70% 上下。加入数据增强和 Dropout 后测试准确率能到 75% 到 80% 这个区间。如果再想往上突破单纯调小模型已经不够通常要换更深或更宽的骨干网络比如 ResNet 一类的结构或者直接迁移预训练模型。配置训练准确率参考测试准确率参考小网络无增强无正则95% 以上70% 左右小网络加增强加 Dropout90% 左右75% 到 80%加深加宽网络调度学习率90% 以上82% 到 85%如果你看到训练准确率已经 95% 而测试准确率只有 70%核心问题不是学习率而是过拟合。此时不要反复调整学习率应该回去检查增强策略和 Dropout。6.2 可视化预测结果与单类精度准确率是一个宏观指标但做彩色图片分类的周任务我建议至少看一眼每个类别单独的表现。CIFAR-10 里最容易混淆的是猫和狗这一类小动物以及汽车和卡车这种外形接近的类别。整体准确率可能还行但某一类召回率会明显偏低。可视化预测结果时用matplotlib很方便import matplotlib.pyplot as plt def show_predictions(model, dataset, device, rows4, cols4): model.eval() figure, axes plt.subplots(rows, cols, figsize(8, 8)) classes dataset.classes for i in range(rows * cols): image, label dataset[i] with torch.no_grad(): pred model(image.unsqueeze(0).to(device)).argmax(1).item() image image.cpu() * torch.tensor((0.2470, 0.2435, 0.2616)).view(3, 1, 1) image image torch.tensor((0.4914, 0.4822, 0.4465)).view(3, 1, 1) image torch.clamp(image, 0, 1).permute(1, 2, 0) row, col divmod(i, cols) axes[row][col].imshow(image) axes[row][col].set_title(fTrue: {classes[label]}, Pred: {classes[pred]}, color(green if label pred else red)) axes[row][col].axis(off) plt.tight_layout() plt.show()上面代码里我手动把归一化后的图像还原成 0 到 1 之间的正常像素这样才能正常显示。查看这类图像时红色标题代表预测错误。如果大量错误集中在某几个类别后续可以朝着类别均衡、难例挖掘的方向优化。6.3 模型保存与后续扩展这周的彩色图片分类任务做到这里已经有了一个可以保存的模型。保存时直接把 state_dict 存下来就够用torch.save(model.state_dict(), color_cnn.pth)后面要推理时先实例化同样结构的模型再load_state_dict把参数加载进去并设置model.eval()。如果再继续往下做可以考虑用预训练模型做迁移学习把输入图片缩放到预训练模型需要的分辨率替换最后的全连接层只训练最后几层。那样通常比从零训练一个随机初始化网络更容易出效果也更能体会到彩色图片分类里“特征复用”的价值。我这次实际做下来最大的感触是彩色图片分类的表面门槛只是多了一个通道真正的门槛却分布在数据管道的每一个细节里。每次遇到结果不理想先把归一化、通道顺序、训练验证模式切换这几件事查一遍能省下大量排查时间。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑