老实说,看到“第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-10:6 万张 32×32 的彩色图片,10 个类别,每个类别 6000 张,其中 5 万张训练、1 万张测试。每张图都很小,用 CPU 也能勉强训练,但真要跑出好看的效果,还是建议开 GPU。
类别列表如下:
| 标签索引 | 类别名 | 标签索引 | 类别名 |
|---|---|---|---|
| 0 | airplane | 5 | dog |
| 1 | automobile | 6 | frog |
| 2 | bird | 7 | horse |
| 3 | cat | 8 | ship |
| 4 | deer | 9 | truck |
它有很多预先做好的 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, padding=4), 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', train=True, download=True, transform=transform_train) test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2, pin_memory=True) test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2, pin_memory=True)这里的顺序是有讲究的: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, padding=4)和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×4=1024。
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_classes=10): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.Conv2d(16, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(2), ) self.classifier = nn.Sequential( nn.Dropout(0.25), nn.Linear(64 * 4 * 4, 256), nn.ReLU(inplace=True), 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(), lr=0.01, momentum=0.9, weight_decay=5e-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 上训练几十步。如果网络本身没有明显 bug,loss 应该能快速降到一个很小的值,训练集准确率可能直接变成 100%。
model = ColorCNN() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.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(f"step {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 我通常固定shuffle=False。
如果发现训练集 loss 已经很低,但测试集准确率还上不去,甚至训练准确率和测试准确率差距越来越大,大概率是过拟合。此时优先做的事是增加数据增强强度,比如再引入随机旋转,或者加一点点随机擦除。其次可以增加 Dropout 比例。不要一开始就盲目加大模型容量,模型越大,在固定数据集上越容易过拟合。
5.4 数据加载速度与进程问题
很多人训练到一半发现每个 epoch 之间卡顿严重,以为模型太复杂,其实瓶颈在数据加载。DataLoader里的num_workers可以设置并行加载进程数,pin_memory=True可以加速数据从内存搬到 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% 左右 |
| 小网络,加增强,加 Dropout | 90% 左右 | 75% 到 80% |
| 加深加宽网络,调度学习率 | 90% 以上 | 82% 到 85% |
如果你看到训练准确率已经 95% 而测试准确率只有 70%,核心问题不是学习率,而是过拟合。此时不要反复调整学习率,应该回去检查增强策略和 Dropout。
6.2 可视化预测结果与单类精度
准确率是一个宏观指标,但做彩色图片分类的周任务,我建议至少看一眼每个类别单独的表现。CIFAR-10 里最容易混淆的是猫和狗这一类小动物,以及汽车和卡车这种外形接近的类别。整体准确率可能还行,但某一类召回率会明显偏低。
可视化预测结果时,用matplotlib很方便:
import matplotlib.pyplot as plt def show_predictions(model, dataset, device, rows=4, cols=4): 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(f"True: {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()。如果再继续往下做,可以考虑用预训练模型做迁移学习,把输入图片缩放到预训练模型需要的分辨率,替换最后的全连接层,只训练最后几层。那样通常比从零训练一个随机初始化网络更容易出效果,也更能体会到彩色图片分类里“特征复用”的价值。
我这次实际做下来最大的感触是:彩色图片分类的表面门槛只是多了一个通道,真正的门槛却分布在数据管道的每一个细节里。每次遇到结果不理想,先把归一化、通道顺序、训练验证模式切换这几件事查一遍,能省下大量排查时间。