简介:面向计算机视觉初学者与深度学习实践者的 VGG16 迁移学习图像分类实战项目,以 62 种大型野外可食用植物为分类对象,解决小型数据集上训练深层 CNN 的难题。压缩包共 2000 个文件,以约 1993 张 JPG 图像为主,另有 4 个 Python 脚本、1 个 txt 说明、1 个 readme 及 1 个 JSON 超参数文件,整体体积 768.42MB。训练采用余弦学习率自动衰减,run_results 目录保存最优权重、训练日志与 loss/accuracy 曲线;预测时运行 predict 即可自动推理 inference 下所有图片,并在左上角绘制前三个概率最高的类别。内置数据集按类别分子目录存放,训练集约 4300 张、测试集约 310 张;若要训练自己的数据,按 readme 规定摆放目录即可,类别数等超参数会自动生成。已有 74 人学习下载,适合课程设计、毕业设计及快速搭建图像分类基线。
1. 62种可食用植物图像分类,为什么用VGG16迁移学习而不是从零训练CNN
62类野外可食用植物的图像分类任务,难的不是“分类网络选哪个”,而是数据量和数据分布。同一株植物在幼叶期和成熟期看起来像两个物种,阴天拍和晴天拍的背景又完全不同。与其从零训练一个深度CNN,更常见的做法是拿ImageNet预训练过的VGG16做迁移学习:把前边的卷积特征直接复用,只重学最后的分类层。这篇内容以这套实战项目的完整落地方案为主线,从结构原理、数据组织、模型替换、训练参数到评估和剪枝都过一遍。适合正在做植物物种识别、林业图像分类或需要快速搭图像分类原型的工程师。
2. VGG16的卷积结构为什么适合做植物图像特征提取器
2.1 16层堆叠的固定模板:卷积块、池化与分类头
VGG16名字里的16,指的是13个卷积层加上3个全连接层。输入图像统一缩放为224×224×3,经过5个最大池化层,特征图尺寸从224降到112、56、28、14,最后变成7×7。通道数从64开始翻倍,到网络深处变成512。这个规律非常规整,因此VGG16的结构图在所有CNN里最好记。
| 阶段 | 操作序列 | 输出特征图尺寸 |
|---|---|---|
| conv1 | 2个3×3卷积(64)+ maxpool | 112×112×64 |
| conv2 | 2个3×3卷积(128)+ maxpool | 56×56×128 |
| conv3 | 3个3×3卷积(256)+ maxpool | 28×28×256 |
| conv4 | 3个3×3卷积(512)+ maxpool | 14×14×512 |
| conv5 | 3个3×3卷积(512)+ maxpool | 7×7×512 |
| classifier | 3个全连接层 | 4096 → 4096 → 1000 |
连续堆叠两个3×3卷积,感受野等效一个5×5卷积,但参数量从25C²降为18C²,还多引入了一次非线性。VGG16的卷积核都极小,结构上没什么花哨的设计,这正是它能被反复复用、容易改造成迁移学习模型的原因。
先用torchvision把模型加载出来看一眼结构,是动手前的标准动作。
import torch from torchvision.models import vgg16, VGG16_Weights model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1) print(model)这段代码以当前torchvision的推荐写法为例。weights=VGG16_Weights.IMAGENET1K_V1会加载在ImageNet上预训练好的权重;如果还在用旧版API里的pretrained=True,也能运行,只是会在新版本里收到废弃提示。打印出来的模型主体由features、avgpool和classifier三部分组成。最后一个卷积层输出的7×7×512=25088个数值,被展平后送入全连接层,这个数字在修改分类头时需要用到。
2.2 预训练权重到底迁移了什么:边缘、纹理与叶片形态
VGG16在ImageNet上训练时,前几层卷积学到的是边缘、颜色和纹理滤波器,中间层学到的是拐角、弧线和重复纹理的组合,最后面的卷积层才偏向物体部件和整体形状。植物分类恰好需要这些特征:叶片锯齿边缘、叶脉走向、花瓣颜色渐变、果实表面质感,都能从预训练权重里找到对应的激活模式。
迁移学习在这里的定位属于直推式迁移学习:源域是ImageNet里的自然图像,目标域是野外植物照片,两者不完全同分布,但低层视觉特征高度重叠。直接复用卷积层参数,只替换最后一层分类输出,是数据量不足时最稳的方案。
| 迁移策略 | 冻结范围 | 训练数据要求 | 收敛速度 | 适用场景 |
|---|---|---|---|---|
| 特征提取 | 冻结features,只训练分类头 | 每类30到50张即可起步 | 快 | 数据少、原型验证 |
| 全量微调 | 全部参数参与训练 | 每类建议至少200张 | 慢 | 数据充足、域偏移大 |
如果62类每类只有几十张,我一般会分两步走:先用特征提取模式把分类头训收敛,再解冻最后两个卷积块,用更小的学习率做微调。直接上全量微调,VGG16有1.38亿参数,训练集不大时很容易把验证集准确率越训越低。
3. 62类植物图像数据集整理与图像增强,让VGG16学到真实特征
3.1 目录结构按ImageFolder整理并保留独立验证集
torchvision里的ImageFolder要求数据按“根目录/类别/图片”的组织方式存放。62个类别,每个类别一个目录,目录名就是类别标签。原始采集的图片文件名往往没有任何语义,先整理成统一结构,后面所有代码都不用动路径。
data/ train/ class_01_荠菜/ IMG_001.jpg IMG_002.jpg class_02_蕨菜/ ... val/ class_01_荠菜/ ...写一个划分脚本,把每个类别的图片按比例拆到train和val。
import os import random import shutil random.seed(42) def split_dataset(src_dir, train_dir, val_dir, val_ratio=0.15): for cls_name in os.listdir(src_dir): cls_path = os.path.join(src_dir, cls_name) images = os.listdir(cls_path) random.shuffle(images) num_val = int(len(images) * val_ratio) os.makedirs(os.path.join(train_dir, cls_name), exist_ok=True) os.makedirs(os.path.join(val_dir, cls_name), exist_ok=True) for img in images[:num_val]: shutil.copy( os.path.join(cls_path, img), os.path.join(val_dir, cls_name, img) ) for img in images[num_val:]: shutil.copy( os.path.join(cls_path, img), os.path.join(train_dir, cls_name, img) )split_dataset("raw_data", "data/train", "data/val", 0.15)执行后,每个类别下会各自生成15%的验证图片。这里的val_ratio=0.15是常用起点,如果总数据量只有每类三四十张,最好提到0.2,保证验证集有足够样本评估。
提示:如果图片是同一台设备在同一地点连续拍摄的,不要用纯随机划分。相邻帧背景和光照几乎一样,会让验证集分数虚高。按拍摄时间或拍摄地点分组后再划分,更贴近真实识别场景。
3.2 训练集图像增强参数:裁剪、翻转与颜色抖动
野外植物照片和ImageNet里的图片差距不小,尤其是光照条件,阴天、树荫、背光都会让同一个物种呈现完全不同的颜色分布。图像增强在这里的作用不是刷数据量,而是让卷积网络忽略这些与物种无关的变化。
from torchvision import transforms mu = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.5, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter( brightness=0.4, contrast=0.4, saturation=0.3, hue=0.05 ), transforms.ToTensor(), transforms.Normalize(mu, std) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mu, std) ])每个增强参数都有对应场景。RandomResizedCrop随机裁剪一部分再缩放,模拟距离变化和枝叶遮挡;RandomRotation(15)容忍拍摄角度倾斜;ColorJitter里brightness=0.4应对强光和逆光,saturation=0.3应对阴天色彩偏淡。归一化的均值和方差使用ImageNet的标准统计量,因为预训练权重是在这个统计条件下学出来的。
验证集只做Resize和CenterCrop,不做随机增强,保证评估结果稳定可重复。Resize到256再CenterCrop成224,是VGG16输入尺寸的常见搭配,直接Resize到224也可以,但前者会让目标稍微放大,识别小叶片时略有优势。
| 增强方式 | 常用范围 | 解决的实际问题 |
|---|---|---|
| RandomResizedCrop | scale=(0.5, 1.0) | 距离远近、遮挡、目标占比变化 |
| RandomRotation | 10到20度 | 拍摄角度不固定 |
| ColorJitter brightness | 0.3到0.5 | 树荫、阴天、背光 |
| ColorJitter saturation | 0.2到0.4 | 季节和湿度导致的颜色差异 |
如果采集到的数据类别严重不均衡,比如某个常见野菜有几千张,某个稀有物种只有二十张,可以用WeightedRandomSampler让每个批次里少数类的出现概率更高。
class_counts = [500, 23, 120, ...] sample_weights = [1.0 / count for count in class_counts] sampler = torch.utils.data.WeightedRandomSampler( sample_weights, num_samples=sum(class_counts), replacement=True )sample_weights按类别样本数的倒数计算,样本越少的类权重越大。num_samples决定了每个epoch总共采样多少张,设为全量数据集的图片总数,会让少数类在一个epoch内被重复采样多次。只要数据集能整体装进内存,这个方案比手工复制少数类图片干净得多。
4. 基于VGG16迁移学习的最小训练代码:特征提取、微调与超参数
4.1 用torchvision加载预训练VGG16并替换最后一层
VGG16的classifier是三段式全连接结构:4096 → 4096 → 1000。前两个4096层在ImageNet上学到的是高层特征的组合方式,可以直接保留。需要替换的只有最后一层,把输出从1000改成62。
import torch import torch.nn as nn from torchvision.models import vgg16, VGG16_Weights def build_model(num_classes=62, finetune=False): model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1) if not finetune: for param in model.features.parameters(): param.requires_grad = False in_features = model.classifier[6].in_features model.classifier[6] = nn.Linear(in_features, num_classes) return modelmodel.features是卷积特征提取部分,finetune=False时这些参数不参与梯度更新,前向传播仍然使用预训练权重。model.classifier[6]指的是classifier这个Sequential里下标为6的模块,也就是最后一层。先取它的in_features再构造新的Linear,避免手写4096这个魔法数字。
有些教程会把整个model.classifier重建为新网络,我一般不建议这么做。除非数据量很大,否则前两个全连接层里有价值的特征组合会被随机初始化的新层覆盖,反而丢失预训练信息。
4.2 训练参数选型:SGD与Adam、学习率、batch size与epochs
迁移学习的训练参数和从零训练不一样。特征提取阶段,只需要训练最后那个Linear层,优化器只接收model.classifier.parameters(),学习率可以给大一点。微调阶段,所有参数都参与训练,VGG16深层卷积的梯度比较敏感,学习率必须调小。
| 训练模式 | 优化器 | 学习率 | batch size | 建议epoch数 |
|---|---|---|---|---|
| 特征提取 | SGD momentum=0.9 | 0.01 | 32 | 20到30 |
| 整体微调 | SGD 或 AdamW | 1e-4 到 3e-4 | 16到32 | 10到20 |
SGD在这个任务上仍然是最稳的选择,收敛曲线比Adam平滑,最终准确率一般也更高。Adam的优势是前期收敛快,适合快速验证代码能不能跑通。weight_decay=1e-4加在优化器参数里,相当于对权重做L2正则,能稍微抑制全连接层的过拟合。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = build_model(num_classes=62, finetune=False).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD( model.classifier.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4 ) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=7, gamma=0.1 )StepLR每7个epoch把学习率乘以0.1,前20个epoch大概会经历两次断层下降。这样做的好处是前期用大学习率快速搜索分类头的参数空间,后期小学习率精细收敛。如果训练曲线在验证集上抖动明显,换成ReduceLROnPlateau按验证loss自动降低学习率更省心。
4.3 训练循环:不用框架也能跑完的最小实现
训练循环本身不复杂,关键点在于把模型切成train和eval两种模式,验证阶段不要计算梯度。
def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss = 0.0 correct = 0 for images, labels in loader: images = images.to(device) labels = labels.to(device) optimizer.zero_grad() logits = model(images) loss = criterion(logits, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) correct += (logits.argmax(dim=1) == labels).sum().item() avg_loss = total_loss / len(loader.dataset) accuracy = correct / len(loader.dataset) return avg_loss, accuracy def evaluate(model, loader, criterion, device): model.eval() total_loss = 0.0 correct = 0 with torch.no_grad(): for images, labels in loader: images = images.to(device) labels = labels.to(device) logits = model(images) loss = criterion(logits, labels) total_loss += loss.item() * images.size(0) correct += (logits.argmax(dim=1) == labels).sum().item() avg_loss = total_loss / len(loader.dataset) accuracy = correct / len(loader.dataset) return avg_loss, accuracylogits.argmax(dim=1)取得每个样本预测概率最大的类别索引,与真实标签比对。损失用CrossEntropyLoss,它内部已经把softmax计算进去了,不需要在网络输出后手动加softmax。images.to(device)把整个batch送到GPU,batch size选择16到32是显存和训练速度的常见平衡点,VGG16参数量大,输入分辨率224×224时,单卡4GB显存跑32的batch基本是上限。
如果数据集不平衡特别严重,可以在构造CrossEntropyLoss时传入weight参数,给少数类更高的损失权重。这个做法和WeightedRandomSampler二选一即可,同时用容易让训练不稳定。
5. 模型评估、过拟合诊断与VGG16剪枝的落地做法
5.1 用classification_report看62类的精度和召回率
整体准确率在植物分类里会掩盖很多问题。62个类别,如果某几个类占了数据量的一半,模型把这几类全认对,整体准确率也能到80%以上,但对其他类可能完全没学会。每次训练完,先跑一个完整验证集的分类报告。
from sklearn.metrics import classification_report import torch model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for images, labels in val_loader: images = images.to(device) logits = model(images) preds = logits.argmax(dim=1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(classification_report( all_labels, all_preds, target_names=class_names, digits=4 ))classification_report会输出每个类别的precision、recall和f1-score。哪个类的recall低,说明这类图片外观变化太大,模型没学到足够的判别特征;哪个类的precision低,说明模型经常把其他类误判成它。野外可食用植物经常出现同科不同种,比如叶片形状接近的菊科植物,这类混淆从准确率曲线上看不出来,只有逐类指标能暴露。
5.2 过拟合的三个信号与早停设置
VGG16容量很大,62类数据量通常不过几千张,过拟合几乎是必然出现的现象,差别只是来得早还是晚。
| 曲线现象 | 判断 | 第一步处理 |
|---|---|---|
| 训练acc继续涨,验证acc开始回落 | 过拟合 | 增强数据增强强度,增大Dropout |
| 训练loss和验证loss都停在高位 | 欠拟合 | 解冻更多卷积层参与微调 |
| 验证acc上下震荡不收敛 | 学习率过大 | 降低学习率或换余弦退火 |
特征提取模式下过拟合主要集中在最后一层Linear,它的参数量是4096×62,约25万个参数,数据量少时很容易记住训练集。最简单的缓解方式是设置早停,只保存验证集上表现最好的那轮权重。
class EarlyStopping: def __init__(self, patience=5, min_delta=0.001): self.patience = patience self.min_delta = min_delta self.best_loss = float("inf") self.counter = 0 def __call__(self, val_loss): if val_loss < self.best_loss - self.min_delta: self.best_loss = val_loss self.counter = 0 else: self.counter += 1 return self.counter >= self.patience每次验证结束后把model.state_dict()保存到文件,如果early_stopping(val_loss)返回True就终止训练。patience=5意味着连续5个epoch验证loss没有明显下降就停止。训练脚本里别忘了在验证集指标更好时覆盖保存权重,而不是等所有epoch跑完再取最后一轮。
5.3 对VGG16进行剪枝:从全连接层开始
VGG16剪枝是常见需求,但要注意剪哪里收益最大。卷积层参数量虽然大,主要在计算量;全连接层的权重矩阵是25088×4096,参数量占比最高,对模型体积影响最直接。
import torch.nn.utils.prune as prune layer = model.classifier[6] prune.l1_unstructured(layer, name="weight", amount=0.3) prune.remove(layer, "weight")l1_unstructured把全连接层中绝对值最小的30%权重置为0,实现稀疏化。prune.remove的作用是固化剪枝掩码,把稀疏后的权重正式写入模型。这里选最后一层剪枝,是因为它本来就是随机初始化的新层,剪完再重训几个epoch就能恢复精度,不会破坏预训练权重。
需要清醒一点:PyTorch默认的非结构化剪枝在GPU上基本没有速度收益,只是减小了参数量和存储体积。要想真正加速推理,要么把权重导出为稀疏格式,要么对通道做结构化剪枝,直接减少卷积输出通道数。对VGG16做通道剪枝复杂度高不少,通常是先剪最后的全连接层看效果,再考虑剪conv5_3这类冗余较明显的卷积层。
6. 进阶应用:给VGG16加一个不认识的拒识阈值
62类植物分类模型训练完成后,会面临一个实际工程问题:模型只有62个输出节点,任何输入都会被分到其中一类,哪怕输入的是一片完全不相干的树叶,模型也会给出一个最高概率的类别。野外识别场景里,“不确定”比“乱认”安全得多。
解决办法是为softmax输出加一个置信度阈值,低于阈值就判定为未知。VGG16的分类头输出的logits经过softmax后,每一类的概率在0到1之间,最大值反映了模型对预测结果的把握。如果最大概率都很低,说明输入不在训练分布内。
import torch import torch.nn.functional as F def infer_with_threshold(model, image_tensor, threshold=0.75): model.eval() with torch.no_grad(): logits = model(image_tensor.unsqueeze(0)) prob = F.softmax(logits, dim=1) max_prob, pred = prob.max(dim=1) if max_prob.item() < threshold: return None, max_prob.item() return pred.item(), max_prob.item()关键参数是threshold。定得太低,未知样本会被漏过去;定得太高,大量真实的可食用植物会被误拒。更稳妥的做法是用验证集标定:收集模型在全部验证集样本上的最大概率,取第5百分位数作为阈值,保证至少95%的已知样本能通过过滤。
all_confidences = [] with torch.no_grad(): for images, _ in val_loader: images = images.to(device) logits = model(images) prob = F.softmax(logits, dim=1) all_confidences.append(prob.max(dim=1).values) confidences = torch.cat(all_confidences) threshold = torch.quantile(confidences, q=0.05).item() print(f"标定阈值: {threshold:.4f}")torch.quantile(confidences, q=0.05)返回排序后第5%位置的值,意味着95%的验证集样本置信度高于这个值。实际部署时如果发现误拒率高,可以把阈值往下调,但随之而来的是未知样本被错分的风险。一个自然的闭环是:把低于阈值的图片单独保存下来,人工标注后归入一个新的未知类别,下次训练时让模型显式学习这个类别,从而持续扩展分类边界。
本文还有配套的精品资源,点击获取