☰
SE注意力机制+VGG16:水果图像分类系统构建实战解析
2026/10/14 5:43:30 网站建设 项目流程

简介:面向水果图像分类场景的深度学习项目压缩包,基于VGG16引入Squeeze-and-Excitation注意力机制构建SE-VGG16-B模型,并通过批量归一化优化训练稳定性,可应对光照变化、背景杂乱等实际条件,适用于农产品品质检测、智能零售等计算机视觉应用。资源共21个文件,以6个Python脚本为核心,覆盖模型定义、训练、预测、混淆矩阵评估与UI交互;另含12张训练过程与结果图表、2个txt说明文档及1个README,整体大小约3.95MB,目录结构清晰便于按需查阅。已有170人学习下载,适合具备一定深度学习基础、希望复现注意力机制分类方案并快速上手的开发者。通过简介与README可了解项目背景、技术路线与实验结论;脚本与图表则具体展示了SE块如何重标定特征通道、BN如何加速收敛,以及模型在水果蔬菜分类、颜色和品种分级上的实际效果,可直接运行复现完整流程。

1. 一个 zip 文件名里的信息量:SE、VGG16 和水果分类

拿到一个叫深度学习_注意力机制_水果图像_分类系统SE_VGG16_B_1741785621.zip的压缩包,先别急着解压。这个命名已经把整个技术栈写清楚了:注意力机制里最经典的 SE(Squeeze-and-Excitation)被塞进了 VGG16 这个骨干网络,用来解决水果图像分类任务。我常年做图像分类工程,看到这个组合的第一反应是:这是一个把「通道注意力」缝进经典卷积网络、在水果类别上做区分的小系统。它要解决的问题很具体——苹果、梨、香蕉、葡萄这些类别,颜色、形状、纹理差异不像人脸那样结构化,模型需要从大量相似特征里挑出真正有区分度的信息,而 SE 正好能帮 VGG16 做到这一点。这篇笔记适合谁?手里有一批水果图片、想从零把分类系统跑起来的人,或者刚接触注意力机制、想知道 SE 到底值不值得加到自己的网络里,加了之后参数怎么调的人。


2. 先弄明白:SE 注意力机制为什么能提升水果分类精度

2.1 SE 模块的计算过程:从全局池化到通道重标定

SE 的全称是 Squeeze-and-Excitation。它不是一个需要重写主干网络的结构,而是一个可以插到任意卷积层后面的小模块。对于输入特征图 U,形状是 H×W×C,SE 先做全局平均池化,把每个通道压缩成一个数值,得到一个长度 C 的向量。这步叫 Squeeze,目的是把空间信息聚合成通道级的统计量,让后续网络能「看见」每个通道整体激活的强弱。

然后是 Excitation。压缩后的向量先通过一个降维全连接层,把通道数从 C 压到 C/r,其中 r 是压缩比,通常取 16;接着过一个 ReLU,再通过一个升维全连接层恢复成 C,最后经过 Sigmoid 输出一组 0 到 1 之间的权重。这组权重表示每个原始通道的重要程度。注意力机制的核心就在这一步:网络不直接修改卷积核,而是学习一个「通道门控」,把重要的通道放大,把不重要的通道抑制。

在水果分类场景里,这个机制很直观。比如红苹果和红番茄颜色接近,但纹理和形状分布不同;如果某个特征通道对应果柄的轮廓,另一个通道对应红色底色,SE 会让模型更依赖前者。下面是一个可以直接用的 PyTorch 实现:

import torch import torch.nn as nn class SEBlock(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y

这段代码的逻辑不复杂。AdaptiveAvgPool2d(1)不关心输入特征图多大,都能池化成 1×1;view(b, c)把池化结果展平成全连接层能吃的形状。两个Linear组成了通道间的瓶颈结构,reduction=16意味着如果输入通道是 512,中间会先压到 32,再升回 512,这样既降低了参数量,又保留了通道之间的非线性关系。return x * y是广播乘法,把权重值乘到每个通道的完整特征图上。

参数上最值得调的就是reduction。16 是论文里的默认值,也是多数工程里表现最稳的起点;如果数据集类别特别少,比如只有 4 类水果,可以试试 8,让瓶颈更宽一点,学到更细的通道关系;但不要低于 4,否则全连接层的参数量会明显上升,小数据集上容易过拟合。

这里还有一个容易被忽略的细节:SE 模块的输出形状和输入完全一致,所以插入到任何卷积组后面都不会破坏特征图尺寸。这也是它为什么能成为工程里最受欢迎的注意力模块之一——你不需要调整前后层的通道数,只需要在合适的位置加一个旁路。很多人在对比实验时问「加 SE 到底加了什么」,其实就是加了一个可学习的通道门控,让网络自己决定朝哪个特征多看几眼。

2.2 为什么选 VGG16:结构简单、插入点明确、迁移学习方便

加注意力机制的骨干网络有很多选择,常见的有 ResNet、MobileNet,但很多工程案例还是会选 VGG16。原因不是它精度最高,而是它的结构足够规整,做实验时能减少变量。VGG16 的features部分是一连串卷积和池化组成的顺序结构,每个阶段都以MaxPool2d作为结束,输出通道分别是 64、128、256、512、512。这就给插入 SE 提供了非常明确的坐标:在每个池化层后面插一个 SE 模块,通道数刚好就是上一层最后一个卷积的输出通道数,不用像 ResNet 那样去拆 BasicBlock 或 Bottleneck。

VGG16 在工业项目里被当作 baseline 用得非常成熟,预训练权重到处都能拿到。迁移到水果分类这种小数据集任务,VGG16 的参数规模是足够的:5 到 10 个水果类别的区分难度不高,把最后一层全连接改成类别数就行。相比 MobileNet,VGG16 在 CPU 上推理确实慢一些,但训练阶段如果你有 GPU,这个代价可以接受。相比 ResNet,VGG16 的梯度传播路径更短,插入 SE 后改动范围可控,对比「加 SE vs 不加 SE」的消融实验时,能很方便地保证只有 SE 这一个变量。

下面这张对比表是我在选择骨干网络时的常用判断依据,你可以直接拿去用于自己的场景:

骨干网络结构复杂度插入 SE 的难度训练显存推理速度适用场景
VGG16低,纯顺序结构低,池化后直接插偏高偏慢小数据集迁移学习、消融实验
ResNet50中,需改残差块中,要定位每一层中等中等中等规模数据、精度优先
MobileNetV3中,含深度可分离卷积中,本身已含 SE 变体低快移动端、边缘设备实时分类
EfficientNet 系列高,宽度深度缩放高,需按原结构改写中等中等数据量大、需要刷点

这张表也解释了为什么很多开源项目里SE_VGG16是最常见的组合之一:它把「注意力机制」这个抽象概念落到了一个结构最简单的主干上,跑通之后再去换 ResNet 或 MobileNet 都容易。反过来说,如果你一上来就在 EfficientNet 里尝试插入自定义注意力模块,光是理解它的 Block 结构就要花不少时间,问题排查成本会成倍上升。

2.3 从命名反推系统结构:B 与时间戳的含义

文件名里的SE_VGG16_B和1741785621不是随机字符。SE_VGG16是网络结构标识,B在我见过的工程里一般表示第二个版本的数据划分或第二次完整训练配置,习惯上 A、B、C 表示不同的实验版本。1741785621是一个 Unix 时间戳,含义是打包这份代码或模型时的时刻。以前我们做实验,模型文件命名经常是model_final.pt,一周后根本分不清它是哪个数据集的产品;后来我学到的习惯是:文件名必须带上网络结构、版本号、精度或时间戳。

这类命名方式在实际协作里非常有用。一份压缩包传到别人手里,光看文件名就能知道:这个系统用的是 SE 注意力机制加 VGG16,版本是 B,打包时间是某个确定的时刻。如果你打算在自己的项目里复现,建议先对照这三个信息确认三件事:第一,你的数据目录结构是不是train/类别名/图片;第二,你的输入尺寸是不是和命名里约定的模型一致,通常是 224;第三,你手里的版本是不是最新的,如果时间戳比你上次训练的模型更早,那可能拿错了包。

时间戳还提醒了一件事:模型和代码是会过期的。你在这份 zip 里看到的时间戳,代表当时的数据分布和实验环境。水果图像分类如果转到了另一个批次的照片,光照、拍摄设备、品类都变了,旧模型不一定还能保持当时的准确率。这也是为什么我强调每次训练都要写清楚数据版本和模型版本,否则几个月后自己都看不懂这个B到底是什么的缩写。命名这件事没有技术含量,但它决定了你后面排查问题要花多少时间。


3. 把水果图像喂进模型:数据整理与预处理

3.1 先编排数据目录:用脚本把原始图片按类别拆开

我见过太多卡在第一步的案例:原始照片放在同一个文件夹里,文件名是照片的编号,唯一的类别线索在某个 Excel 表格里。这种数据不能直接交给 PyTorch 的ImageFolder,因为ImageFolder要求目录结构必须是train/类别名/图片.jpg。所以第一步先把数据整理成标准格式。

常见的整理方式是用 Python 脚本遍历原始目录,按文件夹名识别类别,随机划分训练集和验证集。下面这段脚本可以直接套用:

import os import random import shutil source_dir = "fruit_raw" # 原始数据目录:fruit_raw/apple/xxx.jpg target_dir = "fruit_data" # 整理后的目录 train_ratio = 0.8 # 训练集比例 80% random.seed(42) # 固定随机种子,保证可复现 for root, dirs, files in os.walk(source_dir): label = os.path.basename(root) if not files: continue train_dir = os.path.join(target_dir, "train", label) val_dir = os.path.join(target_dir, "val", label) os.makedirs(train_dir, exist_ok=True) os.makedirs(val_dir, exist_ok=True) random.shuffle(files) split_idx = int(len(files) * train_ratio) for f in files[:split_idx]: shutil.copy2(os.path.join(root, f), os.path.join(train_dir, f)) for f in files[split_idx:]: shutil.copy2(os.path.join(root, f), os.path.join(val_dir, f))

这段脚本的逻辑是:os.walk会递归所有子目录,label取的是当前目录名,所以保证每个类别一个子目录就行。random.shuffle(files)在划分前打乱文件顺序,避免按文件名排序导致某类图片全部进了训练集或验证集。用copy2而不是move,是为了万一划分比例不满意,还能重新执行,不会弄丢原始图片。

train_ratio这个参数很敏感。水果数据通常不会特别多,如果每类只有一两百张,80% 训练、20% 验证意味着验证集每类只有 30 到 40 张。这个规模对评估模型的稳定性是不够的,此时可以把比例改成 0.85 或 0.9。更稳妥的做法是先统计每个类别的图片数量,如果某些类不足 100 张,要么想办法补数据,要么做数据增强,否则后面的训练只能叫过拟合演示,不能叫分类系统。

实际工程里,数据质量检查比划分更重要。我会把每类图片随机抽 5 张拼成一张网格图,肉眼确认有没有错图、有没有没截好的水果、有没有拍到的是包装袋而不是水果本身。这一步看起来原始,但能避免训练完才发现某个类别混入了大量噪声。错误标签是这类小数据集项目里最隐蔽的坑,代码和模型都不会报错,只是准确率永远上不去。

3.2 用 ImageFolder 加载数据集:transform 配置

数据整理好后,加载方式基本固定。torchvision.datasets.ImageFolder会按照子目录名自动生成类别索引,并且利用transform对图片做预处理。这里需要注意训练集和验证集的 transform 不应该一样:训练集需要随机增强,验证集只需要保证输入尺寸和训练一致。

下面是我在这个项目里常用的配置:

import torch from torchvision import datasets, transforms INPUT_SIZE = 224 train_transform = transforms.Compose([ transforms.RandomResizedCrop(INPUT_SIZE, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(INPUT_SIZE), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) train_dataset = datasets.ImageFolder("fruit_data/train", transform=train_transform) val_dataset = datasets.ImageFolder("fruit_data/val", transform=val_transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True )

每个参数都有存在的必要。RandomResizedCrop会随机裁剪一部分区域再缩放到 224 像素,让模型适应水果在不同距离、不同构图下的样子;scale=(0.7, 1.0)的意思是裁掉原图的 30% 以内,对于水果这类主体占比较大的照片,这个下界比较合适,如果缩得太狠,模型可能学到的是「一片红」而不是「一个苹果」。RandomRotation(15)允许图片旋转正负 15 度,超过这个角度会让水果形状失真。ColorJitter对水果分类尤其重要,因为光照变化会导致同一品种的颜色偏移,适度扰动能让模型不过度依赖绝对颜色。

验证集里用Resize(256)再CenterCrop(224)而不是直接Resize(224),这是迁移学习任务里的标准做法。先把短边放大到 256,保留更多上下文,再从中心裁出 224,能减少边界噪声对验证指标的影响。归一化的 mean 和 std 用的是 ImageNet 的统计值,因为 VGG16 的预训练权重是在 ImageNet 上训练的,沿用这套均值方差才能让输入分布和预训练阶段保持一致。这一点新手最容易忽略,很多人的模型不收敛,第一步就是归一化的 mean、std 用了别人随便写的值。

num_workers=4表示用 4 个子进程加载图片。水果图片本身不大,大多数瓶颈在磁盘 IO 而不是 CPU 解码,所以这个值不必盲目调高;Windows 上如果代码if __name__ == "__main__"没写对,num_workers 大于 0 会直接报错,第一次跑通可以先设成 0。pin_memory=True在 GPU 训练时能减少数据从内存到显存的拷贝时间,如果你是 CPU 训练,这个参数开了也没关系,只是没有收益。

3.3 类别不均衡时的后悔药:重采样怎么加

水果拍照时的现实情况是:苹果、香蕉容易收集,杨桃、山竹可能只有几十张。如果不处理类别不均衡,网络会倾向于把不确定的样本都预测成样本量大的类别。常见做法有三个:第一是收集更多数据,这个最直接但成本高;第二是用WeightedRandomSampler让采样器按类别权重抽样;第三是在 loss 里给稀少类别加权。我一般先用第二个,因为它不需要改模型和 loss,改动最小。

from torch.utils.data import WeightedRandomSampler targets = train_dataset.targets class_counts = [targets.count(c) for c in range(len(train_dataset.classes))] total = sum(class_counts) weights = [total / c for c in class_counts] sample_weights = [weights[t] for t in targets] sampler = WeightedRandomSampler( sample_weights, num_samples=len(train_dataset), replacement=True ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=32, sampler=sampler, num_workers=4 )

注意,用了sampler之后,DataLoader里的shuffle参数就必须删掉或者设成 False,因为采样器本身已经完成了打乱。weights = [total / c for c in class_counts]的意思是:某类样本数越少,每个样本被抽中的权重越大;总样本数不变,replacement=True允许重复抽取,让稀少类在同一个 epoch 里多出现几次。这个方案也有代价:数据量少的类别被反复抽,模型容易在那些样本上死记硬背。所以我更建议把它当成一个临时补救方案,核心还是要让每个类别至少有 200 张以上干净图片。

如果不想用重采样,可以改用数据增强来增加稀少类的变化数量,比如对杨桃做更强的翻转、旋转和颜色扰动。但增强只能增加样本的表面变化,不能替代真实场景覆盖。水果分类里,同一个水果在不同光线、不同成熟度下差别很大,真正能救回模型泛化能力的还是尽量多样的实拍数据。类别不均衡是数据工程的典型问题,不要指望注意力机制能替你解决——SE 帮的是特征选择,不是数据分布。


4. 搭建 SE-VGG16 分类系统:从模型定义到训练脚本

4.1 定义 SE 模块并把它插进 VGG16 的 features

第 2 章给出的SEBlock是通用模块,这里要把它真正变成SE_VGG16。在 PyTorch 里,torchvision.models.vgg16返回的模型包含features和classifier两个部分。features是若干卷积、ReLU、MaxPool 组成的顺序结构,classifier 是三个全连接。我们可以在features里遇到MaxPool2d后插入一个SEBlock,通道数取该池化层之前最后一个卷积的输出通道。

import torch import torch.nn as nn from torchvision import models def make_se_vgg16(num_classes=5, pretrained=True): model = models.vgg16( weights=models.VGG16_Weights.IMAGENET1K_V1 if pretrained else None ) new_features = nn.Sequential() last_channels = None for layer in model.features: new_features.append(layer) if isinstance(layer, nn.Conv2d): last_channels = layer.out_channels elif isinstance(layer, nn.MaxPool2d): if last_channels is not None: new_features.append(SEBlock(last_channels)) model.features = new_features # 替换分类器最后一层,适配自己的类别数 num_features = model.classifier[-1].in_features model.classifier[-1] = nn.Linear(num_features, num_classes) return model

这个写法有几个关键点。model.features原本是nn.Sequential,我们用一个新的new_features重新组装:遇到卷积层就正常加入,记录这一层的输出通道;遇到池化层就先加入池化,再把 SE 跟在后面。这样一共会插入 5 个 SE 模块,分别对应 VGG16 的 5 个池化阶段,通道数是 64、128、256、512、512。SE 模块放在池化后面比放在池化前面更好理解,因为池化会缩小空间尺寸,通道重要性可以在低分辨率特征图上计算,计算量更小。

pretrained=True表示加载在 ImageNet 上训练好的权重。这个参数对水果分类这种中小型数据集很重要,能省掉大量训练时间,也能让模型在一开始就具备不错的边缘、纹理、颜色特征。注意我们是在加载预训练权重之后再去修改features的,所以模型里所有原始层都保留了预训练值,只有新增的 5 个 SE 模块需要从头训练。

直接把 SE 插到池化后面,有一个不容易发现的细节:每次MaxPool2d后,特征图的宽高减半,但通道数不变。SE 在这里做的是全局池化,也就是把每个通道在空间上先压成一个数,所以即便特征图只有 7×7,也能稳定算出一个通道级的重要性。如果你尝试把 SE 放在卷积之后、池化之前,理论上也成立,但空间尺寸更大,全局池化的计算成本会高一些,而且通道重要性在更大分辨率上算出来并不一定更准。工程上我更倾向于池化后插入,这也是大多数开源 SE-VGG16 实现选择的插入点。

4.2 微调策略:冻结哪些层,训练哪些层

模型搭好以后,不是所有参数都需要梯度。水果数据集如果只有几千张,从头更新全部参数容易过拟合;常见做法是冻结 VGG16 前几个卷积块,只训练后面的卷积块、SE 模块和分类器。下面这段代码可以直接放在模型定义之后:

for name, param in model.named_parameters(): if "features" in name and "se_" not in name: param.requires_grad = False

这里的意思是:features里除了以se_开头的参数外,全部冻结。SE 模块的名字带有se_,所以不会被冻结;VGG16 原有卷积层的参数保持预训练状态。这样做的好处是训练参数量大幅减少,显存占用也更低,不容易跑偏。如果你的水果数据超过两万张,或者类别之间差异非常小,比如要区分不同品种的苹果,那建议只冻结前两个卷积块,后面全部微调,让高层特征能适应你的具体任务。

值得强调的是,冻结只是让这些层不更新,不代表它们不参与计算。前向传播时,卷积层仍然会执行,只是反向传播时不会计算和存储梯度。所以显存节省不是按照「层数比例」来的,而是取决于你保留了多少激活值和梯度中间结果。如果训练数据量很小,冻结前几个卷积块还能减少模型对训练集的记忆程度,相当于给网络加了一个约束条件。

有一点要特别说明:model.classifier是全连接层,名字是classifier.0、classifier.3、classifier.6这种,上面这段代码不会冻结它,所以分类器会正常训练。分类器参数量不小,但它在迁移学习中本来就是需要重点更新的部分,保留梯度没问题。如果你发现训练损失下降很快但验证准确率上不去,可以试着连classifier里的 Dropout 都检查一下,确认它没有在验证阶段被错误地关闭。

4.3 训练主循环:优化器、学习率、AMP 与模型保存

训练脚本是整套系统的骨架。我用的是 SGD 优化器而不是 Adam,原因是在微调 VGG 系模型时,SGD + momentum 的收敛轨迹更稳,不容易出现后期在验证集上震荡。学习率我通常从 0.005 开始,配合余弦退火衰减;如果从头开始训练,可以提高到 0.01,但迁移学习场景下不要太大。损失函数用CrossEntropyLoss,分类任务的默认选择。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = make_se_vgg16(num_classes=len(train_dataset.classes), pretrained=True).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD( model.parameters(), lr=0.005, momentum=0.9, weight_decay=5e-4 ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30) best_acc = 0.0 epochs = 30 for epoch in range(epochs): model.train() running_loss = 0.0 for images, labels in train_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) model.eval() correct = 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() val_acc = correct / len(val_dataset) if val_acc > best_acc: best_acc = val_acc torch.save({ "model": model.state_dict(), "classes": train_dataset.classes, "best_acc": best_acc, }, f"se_vgg16_b_best_{best_acc:.4f}.pt") scheduler.step() print(f"epoch {epoch+1:02d} | loss {running_loss/len(train_dataset):.4f} | val_acc {val_acc:.4f}")

len(val_dataset)在验证集上取出总样本数,和correct相除得到准确率。preds = outputs.argmax(dim=1)是取每个样本在所有类别上的最大 logits 下标,不需要先 softmax,因为 argmax 结果一样。CosineAnnealingLR的T_max=30要和 epochs 一致,意思是 30 个 epoch 内学习率从 0.005 余弦衰减到接近 0。

模型保存这里有两个细节值得学。第一,只保存state_dict而不是整个模型,这样文件小、跨环境加载更安全;第二,额外保存train_dataset.classes,这个列表就是类别名到下标的映射。后面第 5 章会提到,推理时如果丢掉这个映射,预测结果是一串数字,根本对不上水果名。

warmup也很值得加:前 3 个 epoch 线性把学习率从 0.0001 提到 0.005,能避免收敛初期梯度太大把预训练权重冲坏。如果你用的是 Adam,学习率可以更低,比如 0.0001 到 0.0003,但整体收敛速度不如 SGD 直观。这段代码没有加 AMP 保持结构简单,如果你的 GPU 显存只有 8G 以下,可以引入torch.cuda.amp,但 VGG16 的 batch 32 在 224 输入下,12G 显存一般够用。

训练脚本跑通后,你会得到一个带精度命名的.pt文件,比如se_vgg16_b_best_0.9520.pt。这个命名方式和压缩包里的SE_VGG16_B一致,下次再看到类似的包,就能明白它大概是怎么训练出来的。


5. 训练 SE-VGG16 的常见问题与排查

5.1 验证集准确率比训练集还高:不是模型厉害,是数据泄漏

现象:训练集准确率稳定在 93%,验证集却跑出 98%,怎么看都不合理。

原因绝大多数出在数据划分上。上一章 3.1 里的脚本是随机划分,如果一堆图片本来来自同一次拍摄的连拍帧,被随机分到了训练和验证两个集合里,验证集里就会混入训练集的「近亲」。模型其实记住了画面背景、光照,而不是水果本身,验证准确率自然虚高。

解决方法是按「拍摄来源」分组:比如每个来源目录作为一个整体,把整个目录分到训练或验证,而不是把目录里的文件打散。检查时最简单的办法是抽 10 张验证集图片,看它们的背景是否和训练集重复。水果图如果是网上下载的,很多来自同一组公开数据,图片本身就在多个网站重复出现,这种情况最好提前做去重,网络爬来的图片里经常有同一张图的不同压缩版本。

5.2 预训练权重加载报 size mismatch

现象:执行model.load_state_dict(torch.load("xxx.pth"))时,报size mismatch for classifier.6.weight一类的错。

原因是模型结构改了最后一层全连接,而保存的权重里最后一层还是 1000 维。解决方法是不要直接 load 整个模型的state_dict,而是先构造模型,手动加载预训练权重时设strict=False,只挑能匹配的层;或者像 4.1 里那样让模型加载预训练权重后再替换分类器。

如果你用的是别人给的完整.pt文件,里面可能还包含优化器状态,加载前先打印一下torch.load返回的 keys,不要无脑load_state_dict。打印结果里如果有optimizer、epoch这样的字段,说明这个文件是训练断点而不是纯模型权重,加载方式也完全不同。对于 SE-VGG16 这种自己搭的结构,我强烈建议在保存模型时就使用 4.3 里的字典格式,把model和classes分开存,能省掉很多这类报错。

5.3 显存溢出:先砍输入尺寸,再砍 batch

现象:12G 显存,batch 设 32,训练第一个 epoch 直接 OOM。

原因不一定是 batch 太大,而是 VGG16 在 224 输入时,前面几层特征图的尺寸分别是 112、56、28,显存占用相当高。如果 batch 从 32 降到 8 仍然不行,说明问题在模型本身。

解决:把输入尺寸从 224 降到 160,或者用冻结前几层的方式减少激活值保存;常见的做法是先用with torch.no_grad()跑一遍验证集,观察实际显存占用,再决定调整方向。AMP 在 VGG16 上收益明显,开启后可以省 20% 到 30% 显存,抵消 batch 减半带来的精度损失。这里有个经验:输入尺寸每减一半,特征图存储量会降到四分之一,效果比把 batch 减半明显得多。除非水果图片里的小目标特别多,否则 160×160 的输入对整体准确率影响很有限。

5.4 加了 SE 反而没提升:检查插入位置和 reduction

现象:同样的数据集,VGG16 原版准确率 94%,加了 SE 后变成 92%。

原因可能是插入位置错了,SE 模块被放在了MaxPool2d前面,导致通道数计算出现偏差;或者reduction=16在只有几千张的水果数据集上让通道关系学不到,权重趋近于全 1,等于没加。

解决:先跑通不含 SE 的 baseline,再逐个 stage 插入 SE,观察每个阶段的影响;reduction从 16 改成 8 或 32 各试一次,比较验证集。SE 提升通常只有 0.5 到 2 个点,如果数据本身已经很好分,看不出差异是正常的,不要期待它每次都能救场。我在实际对比时还会做一件事:把训练好的 SE 模块权重打印出来,看 5 个 SE 的输出均值是否接近 1。如果某个 SE 全输出接近 1,说明这个阶段不需要注意力,可以去掉它以省计算量。

5.5 推理时预测结果全错位:忘了保存类别映射

现象:训练时准确率 94%,自己拿一张苹果图测试,脚本返回 3,你对了一下类别列表发现 3 是香蕉。

原因:ImageFolder的类别顺序是按文件夹名的字典序自动生成的,不是按你建立文件夹的顺序。如果你在推理时用了另一份类别列表,或者直接使用model.train_classes的下标,映射就错位了。

解决:保存模型时把train_dataset.classes一起存进去,推理时读取同一个列表。我在 4.3 的训练脚本里已经把这个字段加进.pt文件了,后面的推理脚本会直接用到。这个坑在训练时通常暴露不出来,只有在部署到实际接口时才会炸。如果你是在某个服务里调用模型,一定要把这个列表固化到配置文件里,不要依赖训练脚本的环境变量。


6. 让 SE-VGG16 从训练脚本走到单图预测:推理与检查

6.1 单图推理脚本:返回类别名和置信度

训练和验证都在跑,最后总要落到单张图片预测。下面是通用的推理函数,注意这里的预处理和训练时的验证集保持一致,只是简化成直接 Resize 到 224:

from PIL import Image import torch import torchvision.transforms as transforms def predict(image_path, model, classes, threshold=0.6): transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) image = Image.open(image_path).convert("RGB") batch = transform(image).unsqueeze(0) model.eval() with torch.no_grad(): logits = model(batch) prob = torch.softmax(logits, dim=1).squeeze(0) idx = int(prob.argmax()) confidence = float(prob[idx]) if confidence < threshold: return "unknown", confidence return classes[idx], confidence

这个脚本的关键是把classes从保存的 checkpoint 里读回来,而不是硬编码。torch.softmax提供了置信度,threshold用来处理「不属于任何已知类别」的输入。水果分类系统上线后,最怕的不是分错,而是把一张桔子照片自信地输出「橙子」。加上阈值后,置信度低时返回 unknown,至少不会给用户一个误导。

6.2 加载模型时记住要带回 classes

从.pt文件恢复模型时,常见做法是:

checkpoint = torch.load("se_vgg16_b_best_0.9520.pt", map_location=device) model = make_se_vgg16(num_classes=len(checkpoint["classes"]), pretrained=False) model.load_state_dict(checkpoint["model"]) classes = checkpoint["classes"]

注意pretrained=False,因为权重已经在 checkpoint 里了,不需要再加载 ImageNet 的预训练参数。这里有个我自己的血泪教训:有一次我只保存了模型参数,classes 列表放在训练脚本里,结果训练脚本被清理后,推理代码完全不知道下标对应的水果名,只能靠猜。从那以后我养成了习惯:凡是分类模型,checkpoint 里必带classes字段,推理脚本优先从 checkpoint 读取类别名,绝不在推理代码里硬写一个列表。

最后再提醒一个验证技巧:用一个真正没参与过训练的照片文件夹做外部测试,不要只盯着验证集指标。验证集就算再严格,也是和训练集同源的,只有拿出一批新拍的水果图,才能看出 SE-VGG16 这个系统在真实场景里的表现。这也是模型交付前最后一道必需工序。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询