资讯动态

用VGGNet实现自定义图像分类:从数据准备到模型训练全流程

发布时间:2026/10/19 14:27:27 来源:尧图企业网站定制
1. 项目背景与核心脉络这个项目标题其实已经把整件事的轮廓说得很清楚了用 VGGNet 做自己数据集的图像分类而且特别强调“一步步完整实现”和“自己图像数据集准备”。作为前一篇博客的细化延续这里就不再重复 VGGNet 的基础原理而是把重点放到从零开始把一套分类流程真正跑通的完整细节。先说这个项目适合谁。如果你已经接触过深度学习的框架跑通了官方的 MNIST 或 CIFAR 示例现在手里正好攒了一批自己的图片想训练一个能识别这些小类别的模型那这个项目的每个环节你都会遇到。我不止一次看到有人卡在“模型代码写好了但数据加载器一直报错”或者“训练完精度挺高一到自己随便找的图就完全乱猜”这种地方根源往往是数据准备环节埋了雷或者训练和推理时的数据预处理不一致。整条技术链路包括这几个核心环节自定义图像数据集的目录组织与数据集划分基于 torchvision 进行图像预处理和数据增强VGGNet 的网络结构搭建或迁移学习模型训练的核心参数配置和训练循环实现模型评估、预测可视化与测试集验证先说结论这套流程我用同样方法在好几个项目里跑过无论是几十张的极小数据集还是上千张的常规小数据集VGG 的 11 层或 16 层变体都能在合理时间内收敛前提是数据预处理和迁移学习的打开方式正确。这篇就是把每一步怎么操作、为什么这样操作、踩过什么坑讲透。2. 数据集的制备与目录设计2.1 原始图像的采集与核验很多人一上来就整理图片其实第一步最该做的是把图片先看一遍把明显有问题的剔除掉。我在项目里遇到过的常见问题包括下载或拍摄的图片里混入了空白图、其他人像合影、带有二维码水印的无关图、不同分辨率和横竖比例差异极大的图。这些脏数据对特征学习的干扰非常大尤其是 VGG 这种对局部纹理和形状极度敏感的卷积网络。建议先统一图像格式为 JPEG图片通道统一为 RGB。灰度图虽然也能训练但如果你准备用 ImageNet 预训练权重做迁移学习网络第一个卷积层输入就是三通道灰度图还得转三通道白白损失细节。分辨率不需要统一到某个固定值因为后续数据加载环节会统一缩放但原始文件本身的清晰度会影响最终效果。2.2 目录结构与数据集划分策略我用过的最可靠的方案是torchvision.datasets.ImageFolder的标准目录编排方式就是按类别建文件夹图片按类放入文件结构如下dataset/ ├── train/ │ ├── class_a/ │ │ ├── img1.jpg │ │ ├── img2.jpg │ └── class_b/ │ └── img101.jpg └── val/ ├── class_a/ │ ├── val_img1.jpg └── class_b/ └── val_img10.jpg划分比例上我一般按 8:2 切分训练集和验证集。如果你数据量本身就不大比如每个类别不到 100 张可以尝试 9:1但验证集至少留出每类 15~20 张不然评估指标波动会很大一个 batch 的偶然误差就能让准确率跳好几个点。顺带说一个关键细节一定不要在数据读取的时候做随机采样划分而是要先划分好文件夹再交给ImageFolder去读取。实际项目里我一开始图省事直接用random_split在内存里切结果验证集的分布和训练集几乎重合评估结果虚高换到真实场景就露馅。用文件夹隔离的方式更稳还能保证每次训练和评估看到的数据是完全一致的。2.3 数据增强与归一化配置注意事项数据增强在自定义小数据集上的重要性怎么强调都不为过。from torchvision import transforms, datasets train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(), 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]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里有两个数据增强细节必须留意。第一验证集不要做随机增强只做 Resize 加 CenterCrop这样才能保证评估结果的稳定性和可复现性。第二VGGNet 在 ImageNet 上预训练时输入的标准化参数就是均值[0.485, 0.456, 0.406]、方差[0.229, 0.224, 0.225]这套数值在绝大多数迁移学习场景下不要改改了你就是在和预训练权重对着干。关于RandomResizedCrop的 scale 参数我从经验上来说scale(0.7, 1.0)比默认的(0.08, 1.0)更适合小数据集。默认值裁剪得太狠很多图片的主体对象可能被裁掉大半小数据集上模型反而学不到有效特征。如果你的分类对象本身在图中就占据较小区域可以考虑用RandomResizedCrop加上RandomRotation(degrees10)但要小心某些旋转后面临方向性的对象会被错误标记。3. VGGNet 模型原理与 PyTorch 实现3.1 VGGNet 的结构特点和设计取舍VGGNet 之所以到今天仍被很多个人项目用作分类基准核心在于它证明了“用小卷积核堆深度”这条路是走得通的。连续多个 3×3 卷积堆叠在感受野上等价于一个更大的卷积核但参数量更少而且中间夹杂的 ReLU 激活让网络具备了更强的非线性拟合能力。最常用的 VGG16 结构包含 13 个卷积层和 3 个全连接层卷积部分分成 5 个 stage每个 stage 后接一次 2×2 的最大池化来缩小特征图尺寸。最后一个池化输出的特征图是 7×7×512经过 Flatten 后变成 25088 维向量再进入全连接层。这里要特别提醒一个容易翻车的点正因为 VGG 前两个全连接层的维度高达 4096模型体积和计算量都非常大。如果直接从头训练哪怕图像分辨率只有 224一张 224×224 的图片也要经过五个池化层逐级缩小前向传播的耗时和显存占用都肉眼可见地增长。这也是我强烈建议你优先用预训练权重而不是从头训练的根本原因。3.2 迁移学习适配自定义分类数量用 PyTorch 加载 VGG16 并适配自己的分类数代码上非常直接。import torchvision.models as models model models.vgg16(pretrainedTrue) for param in model.features.parameters(): param.requires_grad False num_features model.classifier[6].in_features model.classifier[6] torch.nn.Linear(num_features, num_classes)这里需要解释两个设计决定背后的逻辑。第一为什么冻结特征提取层只训练分类器。预训练权重在 ImageNet 上已经学会了非常通用的边缘、纹理、形状特征这些特征对绝大多数视觉分类任务都是有效的。小数据集上你只需要让最后的全连接层学会如何组合这些特征这比从头训练要稳定得多收敛也快得多。我实测在几千张图的场景下微调分类器只需要几十个 epoch 就能达到一个不错的效果。第二为什么只替换最后一个全连接层。model.classifier是包含三个全连接层的 Sequential 结构索引 6 对应的就是最后那个输出维度为 1000 的层。替换成自己的分类数即可。中间那两个 4096 的全连接层保留预训练参数它们相当于特征融合器不需要改动。有的朋友可能会考虑把输入分辨率提高来获取更多细节。从操作可行性的角度对 VGGNet 这种全卷积部分加全连接的网络224 是标准输入超过 224 的话你可以把最后一个全连接层的输入维度重新 Flatten 后调整但计算量暴涨小数据集上往往得不偿失。我自己试过 256 输入训练时间几乎翻倍最终准确率提升却非常有限。3.3 训练配置的选择依据PyTorch 中对应 VGG16 的实例化如果你显存吃紧还可以考虑vgg11或者vgg16_bn。批量归一化层能加速收敛但显存占用也略增。import torch import torch.nn as nn device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.classifier.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5)关于优化器的选择我提一个有参考价值的细节当你的训练数据量在每类几百到一千张这个区间时在迁移学习框架下用 Adam 配合较低学习率通常比 SGD 更容易在短时间内达到平滑收敛。因为整体参数量小、训练数据少动量的精细调参带来的收益远不如 SGD 在大数据场景下明显。VGG16 在 ImageNet 上的参数量超过 1 亿直接在全量参数上做反传显存消耗非常夸张。我见过有人用 8G 显存的卡跑 VGG16 的完整微调batch size 只能压到 8 左右结果批归一化统计量都不稳定。冻结特征层之后参与训练的参数骤减到最后一层线性层和可能微调的分类器显存压力小得多batch size 起步 32 没有问题。4. 训练全流程实现与关键参数剖析4.1 数据加载器的实现细节用ImageFolder加载我们前面构建好的目录结构配合DataLoader提供批量读取。from torch.utils.data import DataLoader train_dataset datasets.ImageFolder(dataset/train, transformtrain_transform) val_dataset datasets.ImageFolder(dataset/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)两个配置参数的作用不要忽略了。shuffleTrue对训练集是必须的这样每个 epoch 都能让模型见到不同顺序的数据避免由于同类数据连续出现导致的梯度更新方向偏差验证集不需要打乱。num_workers在 Windows 上如果设置大于 0 偶尔会报错在主流环境或者容器里一般没有问题如果内存紧张就调回 0。pin_memoryTrue对 GPU 训练几乎是无脑开的它把数据预先锁定到页锁定内存数据从主机到设备的拷贝效率更高。但要注意它需要足够的物理内存支撑如果你的机器只有不到 8G 内存而数据吞吐量又大也可以关掉。4.2 训练循环代码的实现训练循环是这套流程的主干我给出的代码每次迭代都会同时计算训练损失和当前 batch 的准确率并在每个 epoch 结束后在验证集上做一次完整评估。from tqdm import tqdm num_epochs 30 best_acc 0.0 for epoch in range(num_epochs): model.train() running_loss 0.0 correct_train 0 total_train 0 for images, labels in tqdm(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) _, predicted torch.max(outputs, 1) total_train labels.size(0) correct_train (predicted labels).sum().item() train_acc 100.0 * correct_train / total_train train_loss running_loss / len(train_dataset) model.eval() correct_val 0 total_val 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total_val labels.size(0) correct_val (predicted labels).sum().item() val_acc 100.0 * correct_val / total_val if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model_vgg16.pth) print(fEpoch {epoch1}: saved best model, val_acc {val_acc:.2f}%) print(fEpoch {epoch1}/{num_epochs} | fTrain Loss: {train_loss:.4f} | fTrain Acc: {train_acc:.2f}% | fVal Acc: {val_acc:.2f}%) scheduler.step()这个循环里有三个问题需要专门讲透。第一训练阶段必须用一个独立的best_acc变量来记录迄今为止最高的验证集准确率并且只在刷新最优记录时保存模型。如果每个 epoch 都保存或者只在最后一个 epoch 保存你会遇到过拟合退化的问题——很多小数据集的模型训练后期在验证集上会回落最后一个 epoch 的模型未必是最优的那个。我复盘过不少训练记录最好的验证集结果经常出现在第 20 个 epoch 左右第 30 个 epoch 时已经有些退化。第二为什么要把验证和训练准确率分开打印而不是只在每 N 个 epoch 打印一次。这是帮助你判断过拟合拐点的最直观途径。当训练准确率持续走高而验证准确率开始持平或掉头就说明模型对训练数据开始“背诵”了理论上应该停止训练或降低学习率。我在实践中发现小数据集上这个拐点出现得远比想象的早比如 VGG16 每类 200 张图往往十几轮就开始出现过拟合迹象单靠看损失曲线是分不清这个区别的。第三model.eval()的调用不是可选项。特别是如果你最终选择vgg16_bn这种带批归一化的结构不切换到 eval 模式的话批归一化层会继续用当前 batch 的统计数据做归一化验证结果会出现令人困惑的抖动和虚高。踩过一次这个坑的人应该都懂。4.3 学习率计划和训练轮次的调参思路我之前给出的StepLR(step_size10, gamma0.5)是一种非常稳的学习率衰减方案每 10 个 epoch 学习率减半让模型在粗调阶段用较高学习率快速下坡在精调阶段用低学习率压实最优区域。实际调参时我的建议是不要一上来就追求一次跑出最好结果先用 10 个 epoch 观察训练趋势。如果 loss 在前两三个 epoch 内没有明显下降大概率是学习率设高了或者数据预处理出了问题如果 loss 降得飞快但在验证集上完全没有相应提升则要考虑是否过拟合开始得太早此时可以调小step_size或者把gamma改成 0.3。训练轮次也不是越多越好。在冻结特征层只训练分类器的场景下30 个 epoch 一般足够。如果做全网络微调那学习率要降到1e-4甚至1e-5同时训练轮次可以适当拉长但风险也随之提高需要配合更频繁的验证集监控。5. 模型测试与推理阶段的核心环节5.1 用测试图像做完整推理训练完成后加载保存的模型权重用一张不在训练和验证集里的独立图片做分类测试。这里的代码其实非常简短from PIL import Image import torchvision.transforms.functional as TF def predict_image(image_path, model, class_names, device, transform): image Image.open(image_path).convert(RGB) image_tensor transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(image_tensor) probabilities torch.nn.functional.softmax(outputs, dim1) top_prob, top_class torch.max(probabilities, dim1) class_name class_names[top_class.item()] confidence top_prob.item() * 100 return class_name, confidence这里class_names对应ImageFolder的class_to_idx映射关系的反向索引在训练时获取在测试时传入。5.2 混淆矩阵与错误样本分析只看单个测试图或者只靠准确率一个指标在小数据集上容易误判模型的真实水平。我每次都会额外输出混淆矩阵来观察模型的偏置情况。import numpy as np from sklearn.metrics import confusion_matrix all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) print(cm)混淆矩阵的价值在于它告诉你模型到底在哪些类别之间容易搞混。比如你的数据里有两个外观高度相似的类它们被互相混淆的频率会明显高于其他配对这直接指出了一个方向该回去补充此类别的训练样本或者考虑在该类别上做更多的针对性数据增强。对于我经历的一个具体项目数据里包含“猫和狗”这类跨品种分类时错误往往发生在特定毛色、特定姿势的图片上用混淆矩阵一眼就能看出来。光靠准确率你是看不出这种局部偏见的。5.3 训练验证与测试三段数据的一致性这一节要讲一个很多项目里最后功亏一篑的细节训练时用的数据预处理流程在预测时必须一字不差地复刻。训练时做了Resize(256)加CenterCrop(224)那预测时对任意一张新图也必须走完全相同的前置缩放和裁剪流程再归一化到同样的均值和方差。如果你训练时用RandomResizedCrop预测时当然不能用它因为随机裁剪没有确定性要换成中心裁剪方式的等价流程保持输入图像的分布与训练时尽量一致。我遇到过一个实际案例训练时始终用 3 通道 RGB 输入到了推理阶段试了一张 PNG 带透明通道的图直接按四通道输入形状对不上报了一堆错。检查源头就是 PIL 读图时没有.convert(RGB)。这类小问题虽不涉及模型逻辑但足够消耗你一下午的排查时间。6. 训练过程中的典型问题与排查经验6.1 损失不下降或剧烈震荡模型训练时 loss 完全不下降或者像锯齿一样来回震荡主要有三种可能性。第一种学习率设置不合理。如果学习率过大梯度更新跨度过大loss 会在高位震荡这时候观察训练集自身的准确率也是起伏不定的。把lr降到1e-4或1e-5重新跑几个 epoch 就能看到平滑的下降曲线。第二种数据本身太随机且增强过强。如果你对每张图都做了非常激进的色彩抖动和随机裁剪某些样本已经被增强得几乎看不出原始类别模型接收的信息过于混乱也会表现为 loss 不降。这种情况下把增强强度调回温和档位比如只保留水平翻转和小范围随机裁剪通常能解决。第三种标签与图像内容存在错配。此前我处理某一份爬取的数据时发现不少图片的分类和内容确实对不上比如夹带了不少与类别无关的图像结果怎么调学习率都是白费。这种问题靠训练曲线看不出来只能靠抽样可视化训练数据来纠察。6.2 GPU 显存不足与训练速度缓慢VGGNet 是一个非常“吃显存”的网络结构即使冻结了特征层前向传播时的中间特征图依然占着显存。显存不足最常见的报错是CUDA out of memory常规手段是先降低batch_size但还有两个不那么直观的操作。一个是关闭梯度计算来降低峰值显存。如果只做验证与预测阶段可以在上下文管理器torch.no_grad()内执行前向传播它不会为中间激活保留额外显存空间。另一个是全网络微调模式下可以考虑开启混合精度训练配合torch.cuda.amp的GradScaler显存占用能直接降到原来的六成左右速度还有提升。如果你发现显存充足但训练速度依然很慢检查num_workers是否太低导致数据加载成为瓶颈。num_workers设为 4 或 8 时CPU 会并行解码和增强图像大幅缩短单个 epoch 的等待时间。6.3 模型泛化能力不足的改善路径验证集准确率始终不理想甚至训练集准确率也不高排除标签错误和代码 bug 后基本就是数据量本身不够模型学。这里有几条经过验证的可行路径数据增强强度拉高一档加入随机旋转、随机尺度变化、遮挡模拟等手段。先冻结特征层训练多轮后再解锁最后的卷积 stage 做低学习率微调。更换更强的预训练骨干但保留 VGG 的分类头逻辑。如果样本实在太少每类不足 50 张优先收集更多数据增强算法解决不了统计意义上的信息不足。其中解锁部分特征层微调是一个非常好的折中方案。操作上很直接for name, param in model.named_parameters(): if features.24 in name or features.26 in name: param.requires_grad True optimizer torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr1e-5 )这样做的意义是底层的通用特征边缘、简单纹理没必要再动而靠近分类器的高层特征可以针对你的数据域做一定适应同时控制了可训练参数量过拟合风险也低。7. 实测效果与心得总结直接说说用这套流程跑一个真实小数据集的表现。类别数是 5每类大约 180 张训练图、40 张验证图用的 VGG16 加冻结特征层训练 30 个 epoch。最终验证集准确率在 87% 到 91% 之间波动不同批次的数据增强随机性会有影响单 epoch 在普通显卡上只需要几秒。整个训练流程流程跑下来包括中间踩坑的时间大概一个下午就完成了。对比方案用同样数据从头训练 VGG16至少需要 60~80 个 epoch 才能达到类似精度而且训练曲线震荡剧烈换一次随机种子结果差异就很大。预训练权重的迁移学习方案在稳定性和时间成本上都有压倒性优势。个人来说我最想强调的还是数据准备的质量。同一套模型代码数据干净规范和混入噪声的乱序排列最终准确率可以相差 20 个百分点以上。很多人执着于调模型结构或学习率却忽略了数据本身才是决定分类性能的上限来源。如果后续想继续扩展这个项目可以考虑的方向包括引入混淆矩阵来做针对性数据补充把分类结果加上 Grad-CAM 可视化判断模型是否在关注图中的合理区域或者更换更强的骨干网络对比效果差异。VGGNet 作为经典架构在这个任务里更多是提供一个稳定扎实的基线但它带来的训练流程和数据管理经验是换到任何其他模型都通用的核心资产。

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

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

免费获取报价 →
↑