1. 为什么选 VGG16 做自定义数据集的图像分类入门1.1 这个项目到底在做什么把 VGG16 这个经典卷积神经网络拿过来在你自己收集的图片数据集上从头训练一遍最终得到一个能对你自己的类别做分类的模型。听起来简单但真正动手做过的人都知道从数据集整理、模型改造、训练调参到最终推理验证中间有一堆细节能把新手卡住。我带过不少刚入门深度学习的同学十个里面有八个在第一次跑自定义数据集时翻车问题基本集中在几个地方数据格式不对、类别映射搞混、最后一层没改、学习率设太大导致loss炸掉、训练集和验证集划分不合理导致过拟合看不出来。这个项目的价值在于它是一条完整的、可复现的流水线。你跟着走一遍就能掌握图像分类任务的核心套路之后换成 ResNet、EfficientNet 或者自己搭一个小网络流程几乎一模一样。适合谁看有基本 Python 基础、了解张量概念、想动手跑通第一个自定义数据集分类任务的开发者。如果你连 PyTorch 都没装过也没关系我会把环境部分写清楚。1.2 为什么偏偏是 VGG16现在网上有很多人会说 VGG16 太老了、参数太多、不如 ResNet。这话对但要看场景。VGG16 的结构极其规整就是卷积堆叠加池化最后接全连接没有任何跳连接、没有复杂的模块设计。对于初学者来说这种规整意味着你调试的时候能清楚地知道每一层在干什么出了问题容易定位。ResNet 的残差连接虽然好但新手看到identity分支和downsample的时候容易懵。另外VGG16 的预训练权重非常成熟即使你的数据集很小用迁移学习的方式微调也能拿到不错的效果。我个人的经验是对于几百到几千张图片、3到10个类别的任务VGG16 微调后的准确率通常能到 85% 以上调得好能上 92%。这个数字对于入门项目来说完全够用。还有一点VGG16 的参数量大约 1.38 亿其中绝大部分集中在全连接层。这意味着你改造网络的时候主要工作就是把最后那几层全连接换掉前面的卷积层基本可以原封不动地拿来当特征提取器。这个改造过程非常直观适合理解“迁移学习”到底在迁移什么。1.3 整体技术路线概览整个项目的流程可以拆成六步环境搭建、数据集准备与加载、模型改造、训练循环编写、训练过程监控与调参、推理验证。每一步都有坑我会在对应章节里详细展开。技术栈方面核心是 PyTorch配合 torchvision 提供的预训练模型和图像变换工具。数据加载用Dataset和DataLoader训练过程用标准的反向传播加优化器更新。可视化部分用 Matplotlib 画 loss 和 accuracy 曲线。这些都是标配不需要额外装什么冷门库。有一点要提前说清楚VGG16 对输入尺寸有要求原始论文用的是 224x224。你自定义数据集的图片尺寸可能五花八门所以预处理阶段必须统一 resize 到 224x224否则全连接层的输入维度对不上直接报错。这个坑我见过太多人踩。2. 环境搭建与数据集准备的关键细节2.1 环境配置版本兼容比什么都重要PyTorch 的版本兼容性是个老生常谈的问题。我建议直接用 conda 建一个独立环境避免和系统里的其他包打架。以下是经过实测的稳定组合conda create -n vgg_custom python3.9 conda activate vgg_custom pip install torch2.0.1 torchvision0.15.2 pip install numpy matplotlib pillow tqdm如果你有 GPU装的时候注意 CUDA 版本要匹配。用torch.cuda.is_available()验证一下返回 True 就说明 GPU 可用。没有 GPU 也能跑只是训练时间会从几分钟变成几十分钟甚至更久取决于数据集大小。提示不要盲目追求最新版本。PyTorch 2.x 之后的 API 有一些变化网上很多老教程用的是 1.x 的写法混着看容易出错。选定一个版本就从头用到尾。2.2 数据集目录结构这一步做对后面省一半事自定义数据集最常见的组织方式是按类别分文件夹。假设你要做一个“猫、狗、鸟”三分类任务目录结构应该是这样dataset/ ├── train/ │ ├── cat/ │ │ ├── 001.jpg │ │ ├── 002.jpg │ │ └── ... │ ├── dog/ │ │ └── ... │ └── bird/ │ └── ... └── val/ ├── cat/ │ └── ... ├── dog/ │ └── ... └── bird/ └── ...这种结构可以直接用torchvision.datasets.ImageFolder加载它会自动根据文件夹名字生成类别标签并且按字母顺序排序。注意是按字母顺序不是按你想要的顺序。所以如果你的类别是“猫狗鸟”实际标签映射可能是 bird0, cat1, dog2。这个映射关系一定要记下来推理的时候要用。我踩过的坑有一次做水果分类文件夹名字用了中文结果ImageFolder读取的时候编码出问题标签全乱了。后来改成英文文件夹名就没事了。所以文件夹名一律用英文小写不要有空格和特殊字符。2.3 数据划分别让验证集骗了你很多人图省事直接从训练集里随机抽 10% 当验证集。这样做不是不行但有两个隐患。第一如果数据集本身类别不平衡随机抽可能导致验证集里某些类别样本极少评估结果不可靠。第二如果同一个物体的多张照片被分到训练集和验证集两边验证准确率会虚高因为模型在训练时“见过”这个物体。正确的做法是先确保每个类别的图片数量大致均衡然后按类别分层抽样。如果同一个类别下有多个不同个体的照片尽量保证同一个体的照片只出现在训练集或验证集中。手动做这件事很累但值得。我一般会写个小脚本按类别遍历每个类别随机抽 20% 放到 val 文件夹。对于小数据集每类少于 500 张验证集比例可以设到 20% 到 30%。数据集大的话10% 到 15% 就够了。2.4 数据增强小数据集的救命稻草自定义数据集通常不大过拟合是最大的敌人。数据增强是成本最低、效果最明显的对抗手段。在 torchvision 里常用的增强操作包括随机裁剪、随机水平翻转、颜色抖动、随机旋转等。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.RandomRotation(15), 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 和中心裁剪。Normalize 用的均值和标准差是 ImageNet 的统计值因为我们要用 VGG16 的预训练权重输入分布最好和预训练时一致。注意RandomResizedCrop的 scale 参数不要设得太激进。我试过设成 (0.3, 1.0)结果有些图片被裁得只剩一个角模型学出来的东西很奇怪。0.7 到 1.0 是比较稳妥的范围。3. VGG16 模型改造与迁移学习策略3.1 加载预训练模型并替换分类头VGG16 在 torchvision 里的调用非常直接import torchvision.models as models import torch.nn as nn model models.vgg16(pretrainedTrue)pretrainedTrue会自动下载 ImageNet 上训练好的权重。下载慢的话可以提前手动下载好放到缓存目录具体路径在~/.cache/torch/hub/checkpoints/下面。接下来是关键的改造步骤。VGG16 的classifier部分是一个 Sequential结构如下(classifier): Sequential( (0): Linear(in_features25088, out_features4096) (1): ReLU(inplaceTrue) (2): Dropout(p0.5) (3): Linear(in_features4096, out_features4096) (4): ReLU(inplaceTrue) (5): Dropout(p0.5) (6): Linear(in_features4096, out_features1000) )最后那个Linear(4096, 1000)是 ImageNet 的 1000 类输出我们要把它换成自己数据集的类别数。假设你的数据集有 5 个类别num_classes 5 model.classifier[6] nn.Linear(4096, num_classes)就这么一行。但这里有个策略选择是只训练最后这一层还是把前面的层也解冻一起训练3.2 冻结与解冻两种迁移学习策略的取舍策略一只训练分类头。把前面所有层的requires_grad设为 False只更新最后那个 Linear 的参数。这样做训练速度极快因为大部分层不需要计算梯度。适合数据集非常小每类几十张的情况因为参数少不容易过拟合。for param in model.features.parameters(): param.requires_grad False for param in model.classifier[:-1].parameters(): param.requires_grad False策略二全部解冻整体微调。所有层都参与训练但用很小的学习率比如 1e-4 或更低。适合数据集中等规模每类几百张以上的情况能让卷积层也适应你的数据分布通常效果更好。我一般的做法是先用策略一跑几轮让分类头先收敛到一个差不多的状态然后解冻所有层用更小的学习率继续训练。这叫“渐进式解冻”实测比一上来就全解冻稳定得多。提示解冻之后学习率一定要降。我见过有人解冻后还用 1e-2 的学习率结果 loss 直接飙到 nan。卷积层的预训练权重是很宝贵的别用大学习率把它冲垮了。3.3 参数计算为什么 VGG16 这么“重”VGG16 的参数量主要集中在全连接层。算一下25088 * 4096 4096 4096 * 4096 4096 4096 * num_classes num_classes。前两项加起来就超过 1 亿了。卷积层的参数量其实很小总共才 1400 多万。这意味着什么意味着如果你只训练分类头实际参与训练的参数量是4096 * num_classes num_classes对于 5 分类来说才 2 万多。训练速度快得飞起但模型容量也受限。如果你解冻所有层参数量瞬间变成 1.38 亿显存占用和训练时间都会大幅上升。显存方面224x224 的输入batch size 设 16 的话VGG16 大概需要 4GB 到 6GB 显存。如果显存不够要么减小 batch size要么用梯度累积模拟大 batch。4. 训练循环编写与关键参数设置4.1 损失函数与优化器选择分类任务用交叉熵损失这是标配criterion nn.CrossEntropyLoss()优化器我推荐用 SGD 加动量或者 Adam。SGD 收敛慢但最终效果通常更好Adam 收敛快但有时候会卡在局部最优。对于 VGG16 微调我一般用 SGD学习率设 1e-3动量 0.9权重衰减 1e-4。optimizer torch.optim.SGD( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3, momentum0.9, weight_decay1e-4 )注意那个filter它只把需要梯度的参数传给优化器。如果你冻结了前面的层但没过滤优化器还是会尝试更新它们虽然梯度是 None 不会真的更新但会浪费计算。学习率调度也很重要。我一般用StepLR每 10 个 epoch 把学习率乘以 0.1scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1)4.2 训练循环的完整写法训练循环的骨架很固定但细节决定成败。下面是我常用的写法加了详细注释import torch from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) num_epochs 30 best_acc 0.0 for epoch in range(num_epochs): # 训练阶段 model.train() running_loss 0.0 correct 0 total 0 train_bar tqdm(train_loader, descfEpoch {epoch1}/{num_epochs} [Train]) for images, labels in train_bar: 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() train_bar.set_postfix(lossloss.item(), acccorrect/total) train_loss running_loss / total train_acc correct / total # 验证阶段 model.eval() val_loss 0.0 val_correct 0 val_total 0 with torch.no_grad(): val_bar tqdm(val_loader, descfEpoch {epoch1}/{num_epochs} [Val]) for images, labels in val_bar: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) val_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) val_total labels.size(0) val_correct (predicted labels).sum().item() val_loss val_loss / val_total val_acc val_correct / val_total scheduler.step() print(fEpoch {epoch1}: Train Loss{train_loss:.4f}, Train Acc{train_acc:.4f}, fVal Loss{val_loss:.4f}, Val Acc{val_acc:.4f}) # 保存最佳模型 if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_vgg16.pth) print(fBest model saved with val acc{best_acc:.4f})几个关键点model.train()和model.eval()必须切换因为 Dropout 和 BatchNorm 在两种模式下的行为不同。验证阶段一定要用torch.no_grad()否则会构建计算图显存爆炸。4.3 学习率与 batch size 的配合学习率和 batch size 是联动的。一般来说batch size 越大梯度估计越准可以用更大的学习率。有个经验公式lr base_lr * batch_size / 256。比如 base_lr 是 0.1batch size 是 32那实际 lr 就是 0.0125。但微调预训练模型时学习率要比从头训练小一到两个数量级。我通常从 1e-3 开始试如果 loss 震荡厉害就降到 1e-4。如果 loss 下降太慢可以适当调大但不要超过 1e-2。batch size 受显存限制。VGG16 在 224x224 输入下batch size 16 大概占 5GB 显存。如果显存只有 4GB可以降到 8同时把学习率也相应调小。注意不要用 batch size 为 1 训练。BatchNorm 在 batch size 太小时统计量不准会导致训练不稳定。最小建议 8。5. 训练过程监控与常见问题排查5.1 用 loss 和 accuracy 曲线判断训练状态训练过程中把每个 epoch 的 loss 和 accuracy 记下来画成曲线能直观看出模型有没有问题。正常的曲线应该是训练 loss 持续下降验证 loss 先下降后趋于平稳或略微上升训练 acc 持续上升验证 acc 先上升后平稳。如果训练 loss 下降但验证 loss 上升说明过拟合了。解决办法加数据增强、加 Dropout、减小模型容量比如冻结更多层、或者早停。如果训练 loss 和验证 loss 都不下降说明欠拟合或者学习率不对。检查学习率是不是太小或者模型是不是根本没在学比如忘了把模型设为 train 模式。如果 loss 出现 nan基本是学习率太大或者数据里有脏数据比如损坏的图片。先把学习率降一个数量级试试不行就检查数据。5.2 常见问题速查表问题现象可能原因解决方法loss 不下降学习率太小、模型未设为 train 模式、数据标签错乱调大学习率、检查 model.train()、重新检查标签映射loss 震荡剧烈学习率太大、batch size 太小降低学习率、增大 batch sizeloss 变 nan学习率过大、数据有异常值、梯度爆炸降低学习率、检查数据、加梯度裁剪验证 acc 远低于训练 acc过拟合加数据增强、加 Dropout、早停、冻结更多层验证 acc 一直不涨学习率太小、模型容量不足、数据划分有问题调大学习率、解冻更多层、重新划分数据显存不足batch size 太大、模型太大减小 batch size、用梯度累积、用混合精度训练速度极慢没用 GPU、数据加载是瓶颈检查 device、增加 num_workers、用 pin_memory5.3 几个我踩过的坑坑一忘了改最后一层。有一次我加载了预训练模型直接开始训练结果 loss 一直不降。查了半天才发现最后一层还是 1000 类输出而我的数据集只有 5 类标签对不上。改完立刻正常。坑二Normalize 的均值和标准差用错。我试过用自己数据集的均值和标准差做 Normalize结果训练效果反而变差了。原因是预训练权重是在 ImageNet 的分布上学的输入分布最好保持一致。后来改回 ImageNet 的统计值就正常了。坑三验证集用了训练集的 transform。训练集的 transform 里有随机裁剪和翻转验证集也用了同样的 transform导致每次验证结果都不一样没法比较。验证集必须用确定性的 transform。坑四num_workers 设太大导致卡死。在 Windows 上num_workers设大于 0 有时候会出问题。如果遇到卡死先设成 0 试试。Linux 上一般设 4 或 8 比较合适。6. 推理验证与模型部署准备6.1 加载保存的模型做单张图片推理训练完之后用保存的最佳模型做推理from PIL import Image # 加载模型结构 model models.vgg16(pretrainedFalse) model.classifier[6] nn.Linear(4096, num_classes) model.load_state_dict(torch.load(best_vgg16.pth)) model model.to(device) model.eval() # 预处理单张图片 img Image.open(test.jpg).convert(RGB) 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]) ]) input_tensor transform(img).unsqueeze(0).to(device) # 推理 with torch.no_grad(): output model(input_tensor) probabilities torch.softmax(output, dim1) predicted_class torch.argmax(probabilities, dim1).item() print(fPredicted class: {predicted_class}, Probability: {probabilities[0][predicted_class]:.4f})注意unsqueeze(0)是给图片加一个 batch 维度因为模型期望的输入是(batch, channel, height, width)。6.2 类别映射的保存与加载前面说过ImageFolder按字母顺序生成标签。训练完之后一定要把class_to_idx保存下来推理的时候才能把预测的索引映射回类别名。import json class_to_idx train_dataset.class_to_idx with open(class_to_idx.json, w) as f: json.dump(class_to_idx, f) # 推理时加载 with open(class_to_idx.json, r) as f: class_to_idx json.load(f) idx_to_class {v: k for k, v in class_to_idx.items()}这个文件看起来不起眼但如果你忘了保存过几天再想推理的时候就得重新猜标签顺序很麻烦。6.3 模型导出与部署的初步考虑如果要把模型部署到生产环境有几个方向可以考虑。一是导出成 ONNX 格式用 ONNX Runtime 推理速度比 PyTorch 原生快一些。二是用 TorchScript 把模型序列化脱离 Python 环境也能跑。三是量化把 FP32 权重转成 INT8模型体积缩小四倍推理速度提升明显但精度会掉一点。# 导出 ONNX dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export(model, dummy_input, vgg16_custom.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})导出 ONNX 的时候注意dynamic_axes参数它允许 batch 维度是动态的这样部署时可以一次推理多张图片。6.4 提升效果的最后几招如果训练完发现准确率不够理想可以试试这几招。第一用更好的数据增强比如 Mixup 或 CutMix它们能显著提升泛化能力。第二用余弦退火学习率调度比 StepLR 更平滑。第三用标签平滑把硬标签变成软标签减少过拟合。第四集成多个模型的预测结果虽然麻烦但效果稳定。我自己最常用的是标签平滑加余弦退火这两个改动很小但通常能带来 1 到 3 个百分点的提升。标签平滑在 PyTorch 里直接用nn.CrossEntropyLoss(label_smoothing0.1)就行余弦退火用torch.optim.lr_scheduler.CosineAnnealingLR。最后再分享一个小技巧训练结束后不要只看最佳 epoch 的模型把最后几个 epoch 的模型也拿来在验证集上评一下有时候最后一个 epoch 的效果反而更好因为学习率已经降得很低了模型收敛得更充分。