ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

基于VGG11与PyTorch的垃圾图像分类实战:从数据准备到模型部署

基于VGG11与PyTorch的垃圾图像分类实战:从数据准备到模型部署 简介基于卷积神经网络VGG的垃圾图像分类项目包面向深度学习初学者、毕业设计及课程设计学生同时也适合需要快速搭建图像分类演示的开发者。项目源自一份高分毕业设计评审成绩95分已通过mac与Windows 10/11环境运行测试可直接作为毕设或课设参考也可在此基础上扩展新功能。压缩包共8个文件主要包含VGG11.py源码、4个训练与测试过程记录csv损失及准确率曲线数据、README部署说明、数据集来源说明以及一个打包好的数据/模型zip整体仅12KB轻量且目录清晰目前已有177人学习下载。该项目的价值在于能够完整还原从数据准备、模型训练到结果评估的流程Python源码注释清晰、csv数据便于直接绘制训练曲线README文档则提供了部署与复现指引适合对照动手练习或作为项目初期的演示蓝本。1. 这个垃圾图像分类项目用 VGG11 在 PyTorch 里跑通了一个完整的毕设链路手头这份资源名为“基于卷积神经网络VGG垃圾图像分类源码部署教程文档全部数据训练好的模型高分项目”压缩包内包含VGG11.py训练脚本、四个 CSV 指标记录文件、训练好的权重、完整数据集以及README.md部署文档。它解决的问题非常具体用卷积神经网络对垃圾图像做多分类且把“数据准备→模型训练→指标记录→预测部署”整条链路都打通了。对计算机视觉入门者和做课程设计的学生来说它的价值不止于 95 分的答辩成绩而是让你看到一份真正能从头跑起来的 PyTorch 代码而非烂大街的 MNIST 教程魔改版。我拆完这份代码后的结论是VGG11 的结构足够简单卷积核堆叠 全连接分类头的设计思想清晰很适合作为第一个“能出指标”的视觉项目。本文从网络结构、数据加载、训练流程、模型调用到部署避坑逐个分析所有命令和参数都基于 PyTorch 环境可以直接照着执行。2. VGG11 结构拆解为什么是 5 段卷积而非 3 段2.1 从 VGG 家族的选型逻辑说起VGG 系列用统一的 3×3 小卷积核堆叠替代大卷积核。VGG11 是这个家族里层数最少的变体配置为64-128-256-512-512 的通道数递增每段卷积后接最大池化最后接三层全连接。做垃圾图像分类时类别之间差异往往在纹理和局部形状上比如塑料瓶的反光边缘、纸张的纤维走向3×3 卷积的感受野堆叠对这类细粒度特征更友好。VGG11 的完整参数布局单位是通道数 conv1_1(64) → conv2_1(128) → conv3_1(256) conv3_2(256) → conv4_1(512) conv4_2(512) → conv5_1(512) conv5_2(512) → FC(4096) → FC(4096) → FC(num_classes)选择 11 层而不是 16 或 19 层的理由很直接垃圾数据集通常只有几千到几万张图像VGG16/19 的参数量会让训练集在几个 epoch 内迅速过拟合验证集损失不再下降而训练集准确率接近 100%。VGG11 在表达能力和参数量之间更平衡单卡 GPU 或纯 CPU 训练时一轮 epoch 的时间也能控制在可接受范围内。项目里的VGG11.py遵循的就是上述结构最后一层全连接输出的num_classes根据实际垃圾类别数修改常见为 6 类玻璃、纸板、金属、纸张、塑料、垃圾。提示从 224×224 输入开始每经过一次最大池化特征图尺寸减半。全连接层接收的输入维度是 7×7×512这个数字在修改输入尺寸时必须同步重算。2.2 卷积核、池化与感受野的配合逻辑代码里每个卷积块都遵循“卷积→批归一化→ReLU→池化”的顺序。批归一化Batch Normalization在这份代码中不是装饰它让每层输入分布稳定学习率可以适度调大而不必担心梯度爆炸。如果你在训练时发现损失曲线在某个数值附近震荡不下降优先检查是不是漏了 BN 层。# VGG11.py 中的卷积块定义逻辑 def make_layers(cfg, batch_normTrue): layers [] in_channels 3 for v in cfg: if v M: layers [nn.MaxPool2d(kernel_size2, stride2)] else: conv2d nn.Conv2d(in_channels, v, kernel_size3, padding1) if batch_norm: layers [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplaceTrue)] else: layers [conv2d, nn.ReLU(inplaceTrue)] in_channels v return nn.Sequential(*layers)这段代码是 VGG11 的核心构建逻辑。cfg是一个列表M表示最大池化层数字代表卷积输出通道数。padding1保证卷积不改变特征图尺寸配合池化的减半行为整套下采样规律是可控的。如果你要改成 VGG13 或 VGG16只需要替换 cfg 列表的长度其他代码不用动。参数层面的选型有一个容易被忽略的点batch_normTrue让模型收敛更快但也意味着推理时 BN 层使用的是训练期间累积的全局均值和方法差。PyTorch 在model.eval()模式下会自动切换如果忘记调用eval()预测结果会出现随机波动。2.3 全连接层的参数规模控制VGG11 的参数量约 1.33 亿其中大部分集中在第一个全连接层FC(512×7×7 → 4096)。垃圾分类数据集不大全连接层最容易过拟合。项目中可以观察到训练集准确率持续上升而测试集准确率停滞就是全连接层在记忆训练样本。缓解手段在代码中加入 Dropout 即可self.classifier nn.Sequential( nn.Linear(512 * 7 * 7, 4096), nn.ReLU(True), nn.Dropout(p0.5), nn.Linear(4096, 4096), nn.ReLU(True), nn.Dropout(p0.5), nn.Linear(4096, num_classes) )p0.5是推荐起点训练阶段随机丢弃一半神经元迫使网络学到冗余特征。预测阶段 Dropout 自动失效不需要额外处理。如果测试集准确率依然不理想优先把Dropout的p调到 0.6而不是急着加数据增强强度。3. 数据组织与预处理ImageFolder 是最省事的方案3.1 数据集目录结构如何组织项目里的全部数据按 ImageFolder 格式组织这是 PyTorchtorchvision.datasets.ImageFolder直接读取的标准布局。目录结构必须严格符合“根目录/类别名/图像文件”的层次类别名即标签名顺序由文件夹名称按字典序决定。dataset/ ├── train/ │ ├── glass/ # 玻璃类图像 │ ├── cardboard/ # 纸板类图像 │ ├── metal/ │ ├── paper/ │ ├── plastic/ │ └── trash/ └── val/ └── (同 train 的子目录结构)训练集和验证集按 8:2 或 9:1 划分类别文件夹名保持一致。ImageFolder会自动为每个类别分配从 0 开始的索引这个索引顺序在加载训练好的模型预测时必须保持一致否则预测结果会张冠李戴。项目中如果发现预测类别总是差一位十有八九是训练和验证时目录里多了一个系统生成的隐藏文件夹比如 macOS 的.DS_Store导致类别顺序错位。3.2 数据加载代码与参数详解# 数据加载与预处理 from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(rootdataset/train, transformtrain_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue)Resize((224, 224))把任意尺寸的输入统一到 VGG 要求的尺寸。RandomHorizontalFlip只做水平翻转因为垃圾图像不存在上下颠倒的语义问题垂直翻转会破坏“瓶子开口朝上”这类方向特征不建议加。ColorJitter的三个参数分别控制亮度、对比度和饱和度的随机扰动范围模拟不同光照环境下的拍摄效果。Normalize用的 mean 和 std 是 ImageNet 数据集的统计值VGG11 预训练权重就是在这个分布下学到的。如果你的数据分布跟 ImageNet 差异极大可以先在数据集上计算真实均值和标准差再替换但绝大多数情况下直接用 ImageNet 参数即可。3.3 Batch Size 与 num_workers 的匹配策略batch_size32在单张 8GB 显存的显卡上正好合适显存不够时改成 16 或 8。注意 batch size 减半时学习率最好也减半否则收敛曲线会明显波动。num_workers4表示用 4 个子进程预取数据Windows 下超过 0 必须放在if __name__ __main__:块里否则会报 RuntimeError。验证集加载时不需要数据增强只做尺寸调整和归一化val_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]) ]) val_dataset datasets.ImageFolder(rootdataset/val, transformval_transform) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4)这里的shuffleFalse很关键因为验证阶段要逐一对比预测标签和真实标签顺序错乱会让准确率计算完全失效。提示垃圾图像分类常见的错误做法是训练和验证共用一套预处理逻辑。验证集加随机翻转会让指标不稳定同一张图每次验证结果可能不同。4. 训练流程与指标记录损失函数、优化器与 CSV 输出4.1 损失函数与优化器的搭配分类任务的标准配置是交叉熵损失nn.CrossEntropyLoss()它内部整合了 LogSoftmax 和 NLLLoss直接输入网络原始输出logits即可不需要在模型最后一层额外加 softmax。优化器选用 SGD 而非 Adam理由和很多视觉任务一致Adam 收敛快但泛化性能往往不如精调过的 SGD。import torch import torch.nn as nn import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model VGG11(num_classes6).to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1)lr0.01是 VGG 系列微调时的常用起点momentum0.9加速收敛并减少震荡。weight_decay5e-4是 L2 正则化的强度有效抑制全连接层的过拟合这个值在 CIFAR-10 和 ImageNet 分类任务里被验证过是较好的区间。StepLR每 10 个 epoch 把学习率乘以 0.1让模型在训练后期用更小的步长在局部最优附近精细搜索。4.2 训练循环中如何记录四种指标项目的 CSV 文件包含四个指标训练集损失、训练集准确率、测试集损失、测试集准确率。每个 epoch 结束后追加一行记录最终可以用 pandas 或 Excel 直接绘制训练曲线。这是毕设答辩时展示模型收敛过程的关键素材。import csv def train_one_epoch(model, loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.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 * 100 return epoch_loss, epoch_acc def evaluate(model, loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) running_loss loss.item() * inputs.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 * 100 return epoch_loss, epoch_acc关键点在model.train()和model.eval()的切换。train()模式下 Dropout 生效、BN 使用当前 batch 的统计量eval()模式下 Dropout 关闭、BN 使用全局统计量。torch.no_grad()在验证阶段关闭梯度计算减少显存占用并加速推理。这两个函数如果漏掉任何一个指标都会失真。torch.max(outputs, 1)返回每行最大值的数值和索引索引即预测类别。loss.item()把标量张量转成 Python 数值累加时乘上inputs.size(0)是因为loss是该 batch 的平均损失需要还原到样本总数维度。4.3 CSV 写入与训练主循环csv_path training_log.csv with open(csv_path, w, newline) as f: writer csv.writer(f) writer.writerow([epoch, train_loss, train_acc, val_loss, val_acc]) epochs 30 for epoch in range(1, epochs 1): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() with open(csv_path, a, newline) as f: writer csv.writer(f) writer.writerow([epoch, round(train_loss, 4), round(train_acc, 2), round(val_loss, 4), round(val_acc, 2)]) print(fEpoch {epoch:02d}: train_loss{train_loss:.4f}, train_acc{train_acc:.2f}%, fval_loss{val_loss:.4f}, val_acc{val_acc:.2f}%) torch.save(model.state_dict(), garbage_vgg11.pth)训练 30 个 epoch 是常见经验值。第 10 个 epoch 学习率下降后验证集损失通常会出现一次明显的进一步下降这是学习率衰减生效的正常信号。如果第 20 个 epoch 后验证集准确率不再提升后续 epoch 可以做早停不必跑满。.state_dict()只保存模型参数体积小且跨设备兼容性更好。4.4 学习率曲线与过拟合识别技巧观察 CSV 里的数值变化能快速判断训练状态现象判断对策训练损失下降但验证损失上升过拟合增大 weight_decay 或 Dropout 的 p训练损失与验证损失都持平不动学习率过小调大 lr 或改用 MultiStepLR 提前衰减验证损失一开始就高于训练损失很多数据分布差异检查验证集是否有异常样本损失曲线剧烈震荡batch size 过小或学习率过大增大 batch size 或减小 lr如果你的 CSV 中训练集准确率在 5 个 epoch 内冲到 90% 以上而验证集准确率只有 60%说明模型容量严重超过数据量。此时优先降低weight_decay不现实正确做法是把 VGG11 的最后一个全连接层输出维度从 4096 降到 1024再配合Dropout(p0.6)观察变化。5. 训练好的模型如何调用于新图像预测5.1 加载权重与新图像预处理项目附带训练好的模型文件部署时不需要重新训练。加载权重前先重新实例化网络结构再调用load_state_dict。结构必须与训练时完全一致包括num_classes否则加载会报 key 不匹配错误。import torch from PIL import Image from torchvision import transforms device torch.device(cuda if torch.cuda.is_available() else cpu) model VGG11(num_classes6).to(device) model.load_state_dict(torch.load(garbage_vgg11.pth, map_locationdevice)) model.eval() # 单张图像预测 preprocess transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(test_plastic.jpg).convert(RGB) input_tensor preprocess(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(input_tensor) probs torch.softmax(outputs, dim1) confidence, pred_idx torch.max(probs, 1) class_names [glass, cardboard, metal, paper, plastic, trash] print(f预测类别: {class_names[pred_idx.item()]}, 置信度: {confidence.item():.4f}).convert(RGB)确保输入图像统一为三通道。如果原始图像是 RGBA 或灰度模式不做转换会导致ToTensor()输出的通道数不是 3输入维度不匹配。.unsqueeze(0)在 batch 维度上扩展因为模型期望输入形状是(batch_size, 3, 224, 224)。5.2 批量预测与结果落盘对一批测试图像做预测时把所有图像放入一个 batch 并行推理比逐张循环快很多。注意DataLoader与ImageFolder的配合即可from torch.utils.data import DataLoader test_dataset datasets.ImageFolder(rootdataset/test, transformpreprocess) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse) preds [] for inputs, _ in test_loader: inputs inputs.to(device) with torch.no_grad(): outputs model(inputs) _, batch_preds torch.max(outputs, 1) preds.extend(batch_preds.cpu().numpy()) # 保存预测结果 with open(predictions.txt, w) as f: for idx, pred in enumerate(preds): f.write(f{test_dataset.imgs[idx][0]}\t{class_names[pred]}\n)test_dataset.imgs是 ImageFolder 内部的样本路径列表索引与preds一一对应。shuffleFalse保证批次顺序不乱否则路径和预测结果会对不上。提示模型训练时类别索引由文件夹排序决定部署阶段如果重新建了 ImageFolder 实例务必确认加载顺序一致。稳妥做法是单独维护一个class_to_idx.json训练时保存、预测时加载。6. 部署避坑与调优技巧环境、显存与数据泄漏6.1 Windows 和 macOS 双平台的兼容性问题项目在 mac/Windows 10/11 上跑通主要踩坑点有三个。第一个是num_workers在 Windows 下必须为 0或训练入口用if __name__ __main__:包裹否则多进程数据加载直接崩溃。第二个是 PyTorch 版本差异旧版 1.x 里torch.load默认加载到 CPU模型张量到 GPU 需要手动.to(device)新版 2.x 行为类似但序列化权重时如果带_metadata信息跨版本加载可能提示不兼容用map_locationcpu加载后再迁移最稳。第三个是 Apple Silicon 芯片的 MPS 加速。device mps可用但和 CUDA 仍有不少算子上限差异训练时建议仍用 CPU 跑小 batch速度不会慢太多。如果你在 Mac 上训练、在 Windows 上部署权重文件.pth是跨平台的直接拷贝即可。6.2 一条 8GB 显存 GPU 可行的显存调优链路VGG11 的显存占用主要在激活值和梯度。batch_size32在 8GB 卡上已经接近临界训练中途显存溢出CUDA out of memory时按以下顺序调整# 第一步降低 batch size 到 16观察显存占用 python train.py --batch-size 16 # 第二步关闭 pin_memory减少锁页内存压力 # 在 DataLoader 中设置 pin_memoryFalse # 第三步开启混合精度训练PyTorch 2.x 内置支持 python train.py --fp16用torch.cuda.memory_summary()可以查看显存分配细节定位是哪一层占用量异常。混合精度训练需要模型中的卷积和全连接层支持 FP16 计算VGG11 的所有算子都兼容。6.3 数据泄漏如何让指标失真README.md中提醒的常见问题之一是数据泄漏表现为验证集准确率出奇得高但实际场景中准确率大幅下降。典型原因是同一张原图出现在训练集和验证集的不同子文件夹里或者通过水平翻转、裁剪生成了近似重复的图像。垃圾图像数据集体积小手工收集时很容易出现重复搬家。检查方法很简单对每个类别算图像文件的 MD5 哈希值比对训练集和验证集是否有重复find dataset -type f -exec md5sum {} \; | sort all_hashes.txt awk {print $1} all_hashes.txt | sort | uniq -d duplicate_hashes.txt如果有相同 MD5 出现在不同路径就是数据泄漏。处理掉重复样本后重新训练测试集准确率可能下降几个百分点但这个数字才是真实可靠的。6.4 精度与速度的最终权衡建议在model.eval()推理阶段可以把torch.no_grad()和torch.inference_mode()替换使用后者在 PyTorch 1.9 中有更少的开销。部署在 CPU 上时开启torch.set_num_threads(4)能充分利用多核但线程数超过物理核心数反而会因上下文切换变慢。如果你打算把模型部署到服务端接口导出 ONNX 格式会让推理速度更快model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, garbage_vgg11.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})dynamic_axes允许 batch 维度动态变化服务端接并发请求时不用固定 batch size。ONNX Runtime 在 CPU 上的推理速度通常比 PyTorch 原生态快 1.5 到 3 倍且不需要依赖完整的 PyTorch 环境生产部署时优先考虑这条路径。至此从网络结构到数据加载、训练记录、模型调用以及跨平台部署的完整链路都已经验证过了剩下的就是根据你的数据集规模微调参数。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进