简介基于VGG16卷积神经网络的六种生活垃圾图像分类迁移学习实战项目面向计算机视觉初学者和需要快速搭建分类方案的开发者包含代码、数据集与预训练权重解压即可运行。压缩包共两千个文件主体为一千九百九十二张图片另有四个脚本、两个文本文件、一份说明文档和一个配置文件整体约七百七十兆。数据集含硬纸板、玻璃、金属、纸张等六个类别训练集有二千零二十四张图片测试集有五百零三张按类别分目录存放项目采用约六千八百万参数的卷积网络配合余弦退火学习率衰减训练十个轮次测试集最高准确率达到百分之八十八并附带最优权重、训练日志及损失精度曲线预测时运行脚本即可自动推理指定文件夹下的图片并标注前三个最可能类别自定义数据只需按说明整理超参数自动生成。目前已有一百三十七人学习可作为迁移学习图像分类的完整实践参考。1. 经典CNN之VGG16生活垃圾图像分类为什么从迁移学习开始做图像分类时最怕的不是模型选错而是手头数据不够。6种生活垃圾图像分类听起来类别少实际收集到的有效样本可能只有几千张光照、拍摄角度、垃圾袋遮挡带来的干扰却一点不少。如果从头训练一个深度CNNVGG16这样的16层网络光参数量就超过1.38亿靠几千张图去拟合结果基本是过拟合训练集准确率99%验证集徘徊在60%。这也是我在实际项目中优先选择VGG16做迁移学习的直接原因ImageNet上预训练好的权重已经让前几层学会了边缘、纹理、颜色块这些通用特征我们只需要把后面跟具体类别相关的全连接层换掉再微调少量参数就能用很少的数据拿到能上线的准确率。VGG16结构规整、预训练权重随处可得、中间层特征可视化方便对生活垃圾这种目标形态差异大但类别内外观多变的场景反而比一味追求网络深度更可控。2. VGG16网络结构拆解与迁移学习的工作方式2.1 从卷积层到全连接层VGG16的骨架VGG16的16指的是13个卷积层加3个全连接层。整个网络被分成5个stage每个stage由若干3x3卷积和max pooling组成。关键的卷积层堆叠方式是用两个或三个3x3卷积串起来模拟更大的感受野而不直接用5x5或7x7卷积。这样做既减少了参数又因为中间多了非线性激活函数让网络表达能力更强。在全连接层部分原始VGG16有fc64096维、fc74096维和fc81000维对应ImageNet的1000类。做迁移学习时前两个全连接层通常保留或至少保留fc6最后那个fc8必须换掉因为我们的输出是6类不是1000类。替换时常见做法是把fc8换成一个nn.Linear(4096, 6)并在前面插入Dropout防止过拟合。网络输入固定是224x224x3的RGB图像。这意味着我们的生活垃圾图片无论原始分辨率多大最终都要被缩放或中心裁剪到这个尺寸。这个固定输入尺寸是VGG16的一个限制但反过来也简化了数据加载时的预处理流程。2.2 冻结与微调迁移学习在VGG16上的两种用法迁移学习在VGG16上有两种常见操作特征提取feature extraction和微调fine-tuning。特征提取的做法是把预训练网络的所有卷积层参数冻结requires_gradFalse只训练新加的全连接层。这种情况下前向传播时卷积层不计算梯度显存占用和训练时间都大幅下降。只要数据集类别特征与ImageNet有重叠比如生活垃圾中的纸张、纸板有类似文本页面的纹理这种做法就能在20轮内让验证准确率冲到85%以上。微调的做法是解冻一部分卷积层通常是解冻最后一个stage的卷积层或者整个网络都参与训练但给卷积层设置一个比全连接层小10倍的学习率。微调适合那些和ImageNet类别差异较大的任务比如彩色垃圾袋、塑料瓶在高光下的反射纹理。我的经验是如果样本量小于5000只微调最后一个stage样本量超过10000再考虑解冻更多层。否则解冻越多层过拟合风险越大。2.3 为什么不用ResNet或EfficientNet模型参数量ImageNet Top-1单张前向耗时CPU迁移学习友好度VGG161.38亿71.5%约200ms结构简单易改易调ResNet502550万76%约80ms残差连接需额外理解EfficientNet-B0530万77.1%约60ms缩放系数多调参复杂从性能上看ResNet50和EfficientNet确实更优但VGG16的卷积层堆叠方式非常规整每个stage的输出shape清晰方便我们绘制特征图、查看中间激活也方便对剪枝做通道分析。对于教学和快速落地的工业项目VGG16是看得懂、改得动的典型。而且迁移学习最看重的不是模型原始精度而是预训练特征的通用性。VGG16的权重大小适中在多个领域被反复验证过出现反直觉的迁移失败概率较低。3. 6种生活垃圾图像数据集的预处理与数据增强3.1 目录组织与Label映射垃圾图像分类数据集的常见组织方式是一个根目录下按类别建六个子目录每一个子目录里放对应类别的图片。比如dataset/glass、dataset/paper、dataset/cardboard、dataset/metal、dataset/plastic、dataset/trash。在PyTorch里用torchvision.datasets.ImageFolder就能自动按子目录名称生成标签。我一般先把数据按8:2划分成训练集和验证集不单独分测试集。垃圾图像分类项目的典型场景是快速验证可行性验证集准确率已经能反映泛化能力。如果后续要上线再从验证集里分一部分出来做测试。标签和类名映射要固定不能因为训练时目录顺序不同导致推理时标签错位。建议在训练脚本里把class_to_idx保存为JSON文件推理时加载。例如import json from torchvision.datasets import ImageFolder train_data ImageFolder(root/data/garbage/train, transformtrain_transform) idx_to_class {v: k for k, v in train_data.class_to_idx.items()} with open(idx_to_class.json, w) as f: json.dump(idx_to_class, f, indent2) print(train_data.classes) print(train_data.class_to_idx)这段代码的含义是ImageFolder会把train目录下的六个子目录按字典序排序每张图片样本的target就是这个顺序对应的索引。保存映射是为了之后用模型对真实场景图片预测时把输出的类别索引转成可读的类别名。3.2 归一化参数要按ImageNet统计值来VGG16预训练模型是在ImageNet上做的标准化训练ImageNet数据集的RGB三个通道的平均值和标准差分别是[0.485, 0.456, 0.406]和[0.229, 0.224, 0.225]。使用预训练权重时输入数据也必须使用相同的归一化参数否则前几层卷积收到的数据分布就是错的迁移效果直接下降。在torchvision.transforms里归一化操作要放在ToTensor()之后from torchvision import transforms normalize transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) val_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])很多初学者在这里犯的错是直接用ToTensor()除以255的归一化结果或者自己算数据集的均值和方差。如果是从头训练模型用数据集自己的统计值没问题但迁移学习场景里必须用预训练模型对应的统计值。计算自己数据集的均值方差会对预训练特征造成偏移相当于把网络输入域改了。3.3 数据增强不是越多越好关键是随机裁剪与水平翻转垃圾图像数据集的难点在于同一个类别里物品形状变化很大比如塑料瓶有透明、绿色、蓝色纸张有白色A4纸、报纸、纸箱碎片。数据增强要围绕模拟同一物体不同视角来设计而不是无脑叠加所有变换。对于VGG16固定224x224输入RandomResizedCrop是效果最好的增强方式它随机裁剪原图的0.6到1.0倍面积并缩放到224能模拟不同距离的拍摄。RandomHorizontalFlip对垃圾图像基本无害因为垃圾物品没有左旋右旋的语义方向问题。ColorJitter可以适当加但亮度扰动幅度不要超过0.3否则塑料和金属的反射特征会被过度改变。不要用旋转90度或随机旋转因为很多垃圾有重力方向感比如塑料瓶和易拉罐旋转90度后可能让纹理方向变得不符合真实场景。也不要加大范围椒盐噪声VGG16的卷积特征本身对噪声不敏感加了反而让预训练特征退化。4. 用PyTorch实现VGG16迁移学习训练与评估4.1 加载预训练模型并替换分类头常见做法是用torchvision.models里自带的VGG16预训练权重加载模型因为该权重是在ImageNet上完整训练过的。替换分类头时注意VGG16的classifier是一个nn.Sequential最后一个线性层的输入维度是4096输出是1000。我们可以直接重建一个classifier或者只替换最后一行。import torch import torch.nn as nn from torchvision import models def get_vgg16_transfer(num_classes6, freeze_backboneTrue): model models.vgg16(weightsmodels.VGG16_Weights.IMAGENET1K_V1) # 把原始的6层分类器替换成适配6类垃圾的3层分类器 model.classifier nn.Sequential( nn.Linear(25088, 4096), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(4096, 1024), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(1024, num_classes), ) if freeze_backbone: for param in model.features.parameters(): param.requires_grad False return model这里有一个关键点model.classifier的输入维度25088对VGG16是固定不变的它由最后一个卷积池化层输出的7x7x512展平而来。如果你在features里额外加了卷积层或改了输入分辨率这个数字就必须重新计算。替换分类头时保留了4096维的中间层是因为预训练好的fc6、fc7权重虽然被丢弃了但4096维的全连接特征对细粒度分类仍然足够。freeze_backbone参数控制是否冻结features。注意我们没有冻结整个classifier因为重构后它是随机初始化的必须参与训练。4.2 冻结特征层与分层设置学习率冻结features后训练时优化器只会更新classifier的参数。但如果后续想解冻最后几个卷积层进行微调需要用一个循环把特定的参数设置为requires_grad True。实际项目中更精细的做法是给不同的参数组设置不同的学习率classifier用1e-3features中靠近输出的层用1e-4或1e-5。def split_parameters(model): classifier_params [] backbone_last_stage_params [] other_backbone_params [] for name, param in model.named_parameters(): if not param.requires_grad: continue if classifier in name: classifier_params.append(param) elif features.28 in name or features.29 in name: # features.28 和 features.29 是最后一个卷积池化层 backbone_last_stage_params.append(param) else: other_backbone_params.append(param) return [ {params: classifier_params, lr: 1e-3}, {params: backbone_last_stage_params, lr: 5e-4}, {params: other_backbone_params, lr: 1e-5}, ]分层设置学习率的原理是靠近输入的层学到的是通用边缘纹理与垃圾还是ImageNet无关学习率要小靠近输出的层已经逐渐抽象出和类别相关的语义特征迁移时更需要调整学习率可以大一些。如果所有层都用同一个学习率要么微调不动导致训练慢要么破坏底层特征。4.3 训练循环与指标监控训练循环的本质是反复做前向计算、求loss、反向传播、优化器更新。对于VGG16由于模型参数多Batch Size设置8到32之间比较常见否则显存溢出。下面是完整的训练函数import torch import torch.nn as nn import torch.optim as optim from tqdm import tqdm def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in tqdm(train_loader, descTraining): 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) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc我习惯在每个epoch后额外跑一个验证循环同时监控验证loss和准确率loss下降通常比准确率更早暴露过拟合。criterion使用nn.CrossEntropyLoss()它内部已经做了softmax和log运算不需要在模型输出后再加softmax。训练时模型输出的是未归一化的logits推理时取argmax就能得到预测类别索引。4.4 混淆矩阵与单张图像推理训练结束后用验证集做一次推理记录所有真实标签和预测标签然后绘制混淆矩阵。垃圾图像分类的混淆矩阵能直观看出哪两类在混淆——常见的困惑对是cardboard纸板和paper纸张因为高光下纸箱和纸片纹理类似。import numpy as np from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt def evaluate_confusion(model, val_loader, device, idx_to_class): model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in tqdm(val_loader, descEvaluating): images, labels images.to(device), labels.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) report classification_report(all_labels, all_preds, target_nameslist(idx_to_class.values())) print(report) plt.figure(figsize(8, 6)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelsidx_to_class.values(), yticklabelsidx_to_class.values()) plt.xlabel(Predicted) plt.ylabel(True) plt.savefig(confusion_matrix.png, dpi150)推理单张图片时要注意把输入也做和验证集一样的预处理Resize到256CenterCrop到224归一化。模型输出后直接torch.max拿到类别索引再用之前保存的idx_to_class.json转成类名。这里最容易出错的是忘记调用model.eval()导致Dropout和BatchNorm在推理时仍然生效输出结果不稳定。5. VGG16垃圾图像分类的3个必调参数与验证技巧5.1 学习率与Batch Size的联动调整VGG16在迁移学习中新加的分类头学习率通常从1e-3起步如果验证准确率在前5个epoch没有明显上升说明学习率偏大需要降到3e-4。解冻的卷积层学习率保持在1e-5到5e-4这个区间。一个容易被忽略的问题是增大Batch Size时梯度的方差减小可以适当提高学习率但VGG16显存瓶颈明显Batch Size往往只能固定一个值。所以建议锁定Batch Size为16只调学习率。5.2 验证集准确率和损失同时监控不能只看验证准确率。垃圾图像分类中类别不均衡很常见比如glass样本少、plastic样本多准确率可能被多数类拉高。我一般在每5个epoch打印一次classification_report看每个类别的precision和recall。如果某个类别recall特别低说明模型对这个类的形态变化覆盖不够就要针对性增加该类的数据增强强度比如对瓶装垃圾加大亮度扰动对纸类加大随机裁剪比例。5.3 用CosineAnnealingLR做学习率退火迁移学习容易在训练后期遇到准确率震荡不收敛。固定学习率训练时最后的收敛阶段需要一个更小的学习率来精细调整。常见做法是scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max30, eta_min1e-6 )把训练总轮数设为T_max学习率会在0到eta_min之间沿余弦曲线下降。这样前期能大步搜索后期小步收敛通常比等间隔衰减更稳定。配合Early Stopping连续10个epoch验证loss不下降就停止训练保存最好的一次权重。最后一个重要验证技巧用训练好的VGG16去跑一张和数据集风格完全不同的垃圾图片比如网上搜一张户外垃圾桶的照片。如果模型输出类别彼此之间置信度很接近说明它没有真正学到区分性特征需要回到数据增强和微调层数上去找原因。这一张图的响应比验证集上多0.1个百分点的准确率更有说服力。本文还有配套的精品资源点击获取