PyTorch轻量级CNN实战：MobileNetV3植物幼苗分类与数据增强

简介面向PyTorch图像分类初学者和深度学习实践者MobileNetV3实战资源以MobileNetV3为模型主干使用12种类别的植物幼苗图像数据演示完整分类流程。资源压缩包共包含2461个文件其中2454张PNG格式植物幼苗图片构成数据集5个Python脚本覆盖模型调用、数据加载、训练验证与预测环节另附1个预训练权重文件和1个PDF说明文档整包总大小约951.83MB。已有1714人学习下载内容可帮助读者掌握从torchvision.models调用MobileNetV3、自定义Dataset加载方式以及Cutout与Mixup两种数据增强的具体实现同时涵盖训练和验证流程、余弦退火调整学习率以及预测的两种写法等关键知识点。整包结构清晰适合用作分类项目参考借助PDF配套说明可快速迁移到其他类似图像分类任务中。1. MobileNetV3 轻量特性与植物幼苗分类任务在温室或农田边缘设备上做植物幼苗识别最头疼的不是模型精度不够而是推理速度与内存占用撑不起常规 ResNet。植物幼苗数据集一共 12 个类别每类图像数量并不算多如果直接训练一个深层卷积网络很容易陷入过拟合而且部署时帧率难看。这里选择 MobileNetV3原因很直接它在 ImageNet 上精度接近 ResNet50但计算量只有后者的约五分之一配合深度可分离卷积和注意力机制小样本分类任务里迁移学习的效果往往比想象中更好。下文将从 torchvision 模型加载、自定义 Dataset、Cutout 与 Mixup 增强、余弦退火训练到预测写法完整拆一遍我用 PyTorch 复现这个项目的全过程适合做嵌入式视觉或刚入门图像分类的读者参考。2. torchvision 加载 MobileNetV3 与自定义 Dataset 实现MobileNetV3 在原论文里分 Large 和 Small 两个版本分别侧重精度和速度。torchvision 里提供了mobilenet_v3_large和mobilenet_v3_small都带预训练权重。这里的植物幼苗分类任务只有 12 个类别输入图像尺寸也不算大所以直接用预训练权重做迁移学习是常规做法主干特征提取能力已经足够好只需要替换最后的全连接层并微调即可。2.1 torchvision 中的 MobileNetV3 模型接口torchvision.models从 0.9 版本开始正式内置 MobileNetV3调用方式非常固定。加载模型时需要注意pretrained参数新版中建议使用weights枚举不过为了代码简洁这里仍然用pretrainedTrue。import torch import torch.nn as nn import torchvision.models as models model models.mobilenet_v3_large(pretrainedTrue) # MobileNetV3 Large 的分类器是 Sequential最后一层是线性层 in_features model.classifier[-1].in_features model.classifier[-1] nn.Linear(in_features, 12) # 如果想冻结主干只训练分类器可以这样设置 for name, param in model.parameters(): if classifier not in name: param.requires_grad False这段代码的逻辑是先拿到预训练模型然后读取classifier[-1]的输入维度替换成输出为 12 的线性层。因为植物幼苗图像与 ImageNet 图像风格接近冻结主干能显著减少训练时间防止小数据集上过早过拟合。如果要进一步微调可以把requires_grad设为 True或者只冻结前几层。models.mobilenet_v3_small的接口完全一样只是classifier结构略有不同但替换最后一层的方式相同。2.2 自定义 Dataset 加载植物幼苗图片PyTorch 的图像分类项目通常不使用ImageFolder直接加载因为数据集的目录结构可能不是标准分类格式或者需要在读取时加入更多预处理。自定义 Dataset 类是更可控的方案。植物幼苗数据集中每一类一个文件夹文件名就是类别名采样时还需要统计每个类别的样本数量。import os from PIL import Image from torch.utils.data import Dataset from torchvision import transforms class SeedlingDataset(Dataset): def __init__(self, root_dir, transformNone): self.classes sorted(os.listdir(root_dir)) self.class_to_idx {cls: i for i, cls in enumerate(self.classes)} self.samples [] for cls in self.classes: cls_dir os.path.join(root_dir, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.samples.append((os.path.join(cls_dir, fname), self.class_to_idx[cls])) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label这里有几个关键点。self.classes通过排序保证每次运行类别顺序一致否则标签会跳动。class_to_idx是类别字符串到整数索引的映射后续预测输出时也要用这个映射还原名称。__getitem__里用Image.open(path).convert(RGB)统一转成三通道避免灰度图或带透明度图片导致张量维度错误。实际项目中如果图片数量很大可以把samples直接存成列表也可以在此基础上重写__getitem__加入缓存逻辑。2.3 模型输入尺寸与预处理参数MobileNetV3 在 torchvision 官方实现中输入尺寸是 224x224但这个尺寸不是强制的。植物幼苗图像中叶片纹理比较细腻过小的输入会抹掉细节过大的输入则让模型计算量翻倍。项目里实际使用 224 是平衡速度和精度的选择。transform 需要与预训练权重匹配否则会严重影响收敛速度。处理步骤参数说明Resize256先放大到 256给随机裁剪留出空间RandomResizedCrop224随机裁剪同时覆盖缩放和裁剪增强RandomHorizontalFlipp0.5水平翻转植物幼苗方向敏感度低可用ToTensor无将 PIL 图像转为 [0,1] 张量Normalizemean[0.485,0.456,0.406], std[0.229,0.224,0.225]ImageNet 统计值预训练模型的输入分布train_transform transforms.Compose([ transforms.Resize(256), transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(), 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]) ])验证集不启用随机增强只做中心裁剪保证评估结果稳定。scale(0.8, 1.0)限制了裁剪面积比例避免裁掉幼苗主要部位。这里把 normalize 放在最后因为 ToTensor 会改变数组顺序并归一化到 0-1Normalize 再按通道减均值除标准差。如果使用的是自己训练的权重没有预训练统计值可以不写 Normalize 或用数据集自身的均值方差。3. Cutout 与 Mixup 数据增强在幼苗图像中的落地植物幼苗数据集样本数量有限12 类幼苗彼此在颜色和叶形上高度相似直接用基础翻转裁剪增强模型在验证集上的表现往往不够稳定。Cutout 随机遮挡图像的一个矩形区域强迫模型关注剩余结构和纹理而不是某个标志性色块。Mixup 则把两张训练图按比例混合同时把标签也做线性插值让模型学到的决策边界更平滑。这两种增强方式都适合小型数据集并且实现成本很低。3.1 Cutout 增强原理与实现Cutout 的核心思想是在输入图上随机挖掉一块方形区域并用固定像素值填充。由于图像已经经过 Normalize填充值通常设为 0相当于该区域所有通道归零。在 PyTorch 里可以继承transforms.RandomCrop或直接写成函数。import random import torch class Cutout: def __init__(self, size, fill_value0.0): self.size size self.fill_value fill_value def __call__(self, img): if self.size 0: return img h, w img.shape[1], img.shape[2] y random.randint(0, h) x random.randint(0, w) y1 max(0, y - self.size // 2) y2 min(h, y self.size // 2) x1 max(0, x - self.size // 2) x2 min(w, x self.size // 2) img[:, y1:y2, x1:x2] self.fill_value return img这里img是形状为 (C, H, W) 的张量。随机生成中心点后计算矩形边界再做越界裁剪。fill_value0.0对应 Normalize 之后的黑色因为 ImageNet 的均值和标准差不是均匀的所以 0 不等于原始图像的黑色但这并不重要网络学习时会把该区域视为无关信息。size是矩形边长一般取输入尺寸的 1/4 到 1/3比如 224 输入用 48 或 64。太小起不到遮挡作用太大会把主体目标全部遮住导致训练不稳定。3.2 Mixup 增强原理与实现Mixup 在 batch 层面操作而不是单张图。它从同一个 batch 中随机打乱样本顺序然后对图像和标签分别做线性插值。标签不再使用 one-hot 硬标签而是两个标签的加权和。对于多分类交叉熵损失我们可以直接使用混合后的软标签。def mixup_data(x, y, alpha0.2): if alpha 0: lam random.betavariate(alpha, alpha) else: lam 1.0 batch_size x.size(0) index torch.randperm(batch_size, devicex.device) mixed_x lam * x (1 - lam) * x[index] y_a, y_b y, y[index] return mixed_x, y_a, y_b, lam def mixup_criterion(criterion, pred, y_a, y_b, lam): return lam * criterion(pred, y_a) (1 - lam) * criterion(pred, y_b)alpha0.2是 Beta 分布的参数取值越小lam 越接近 0 或 1即两张图混合程度较轻。torch.randperm生成一个随机排列索引将当前 batch 与自己的打乱版本混合。mixup_criterion对两个标签的交叉熵损失做加权平均总梯度也等于两部分损失的加权和所以训练过程等价于同时学习两个新样本。Mixup 的副作用是训练损失不会降到非常低因为标签本身是概率混合这与 Cutout 不同评估时不需要做任何特殊处理验证集仍然用标准交叉熵。3.3 两种增强的组合顺序与训练集成Cutout 属于空间增强Mixup 属于样本间增强两者可以串起来。实际操作时Cutout 放在 transform 里对单张图操作Mixup 放在 DataLoader 返回 batch 之后训练循环内部执行。组合方式如下。train_transform transforms.Compose([ transforms.Resize(256), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), Cutout(size56) ]) for batch_idx, (inputs, targets) in enumerate(train_loader): inputs, targets_a, targets_b, lam mixup_data(inputs, targets, alpha0.2) outputs model(inputs) loss mixup_criterion(criterion, outputs, targets_a, targets_b, lam) optimizer.zero_grad() loss.backward() optimizer.step()这里有个顺序问题Cutout 在 Normalize 之后做这样输入张量已经是 float 类型且经过标准化随机填充 0 不影响梯度计算。如果放在 ToTensor 之前做需要把 PIL 图像转换成 numpy 数组并用像素值 128 填充相比之下张量操作更简洁。Mixup 里的lam会随 batch 变化每个 batch 的混合比例都不同这本身就带了随机性。需要注意mixup_data中lam是标量广播到整个 batch所以同 batch 内所有样本混合比例相同。如果想要更细粒度可以按每张图生成单独的 lam不过原版 Mixup 论文用的是统一 lam这里沿用原版即可。4. 余弦退火与训练验证流程参数详解网络结构和数据增强确定后影响最终精度的就是优化器和学习率策略。固定学习率在训练后期容易在局部最优附近震荡而余弦退火能在每个周期内先大后小地调整学习率使得模型快速跳过鞍部并精细收敛。这里结合训练循环、验证循环和参数设置展示一次完整训练过程。4.1 余弦退火学习率调整器PyTorch 的CosineAnnealingLR根据公式lr lr_min 0.5 * (lr_max - lr_min) * (1 cos(pi * T_cur / T_max))调整学习率。其中T_max是半个周期长度通常设为epochs这样从初始学习率余弦下降到最小值。import torch.optim as optim from torch.optim import lr_scheduler optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) scheduler lr_scheduler.CosineAnnealingLR(optimizer, T_max30, eta_min1e-5) # 每个 epoch 结束后调用 scheduler.step()这里使用 SGD 而不是 Adam原因是预训练模型微调时SGD momentum 的泛化能力通常更好尤其是在小数据集上。lr0.01是针对全模型微调的量级如果只训练分类器可以适当调大到 0.05。eta_min是学习率下限设为 1e-5 防止后期完全停止更新。T_max30表示 30 个 epoch 完成一个完整余弦周期。CosineAnnealingWarmRestarts也可以考虑它对较长训练次数更友好但植物幼苗分类任务 30-50 epoch 足够简单余弦更稳定。4.2 训练循环与验证循环训练循环里Mixup 只作用于训练集验证集保持原始数据。每轮结束时用验证集计算准确率和损失以此决定是否保存模型。def train_one_epoch(model, loader, optimizer, criterion, use_mixupTrue): model.train() total_loss 0.0 correct 0 total 0 for inputs, targets in loader: inputs, targets inputs.cuda(), targets.cuda() if use_mixup: inputs, targets_a, targets_b, lam mixup_data(inputs, targets, alpha0.2) outputs model(inputs) loss mixup_criterion(criterion, outputs, targets_a, targets_b, lam) else: outputs model(inputs) loss criterion(outputs, targets) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * inputs.size(0) if use_mixup: _, preds outputs.max(1) correct (lam * preds.eq(targets_a).sum().item() (1 - lam) * preds.eq(targets_b).sum().item()) else: preds outputs.argmax(1) correct preds.eq(targets).sum().item() total targets.size(0) return total_loss / total, correct / total验证函数与训练类似但需要取消 Mixup并且用torch.no_grad()关闭梯度计算。def validate(model, loader, criterion): model.eval() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for inputs, targets in loader: inputs, targets inputs.cuda(), targets.cuda() outputs model(inputs) loss criterion(outputs, targets) total_loss loss.item() * inputs.size(0) preds outputs.argmax(1) correct preds.eq(targets).sum().item() total targets.size(0) return total_loss / total, correct / total训练时计算correct在 Mixup 情况下有两种思路一是忽略混合标签直接按预测类别跟两个原始标签比对再乘以各自权重二是只统计正常训练时不开启 Mixup 的准确率。这里采用第一种因为它能反映模型在混合数据上的判断一致性。验证时model.eval()会关闭 Dropout 和 BatchNorm 的 batch 统计这是必须的否则验证结果会随机波动。4.3 训练参数设置表与调参经验下面这套参数是基于 224x224 输入、12 类幼苗数据、约几千张训练图的常见配置可以直接作为起点。参数名值说明batch_size32显存不够可降到 16但要同时调低 learning rateepochs30加上余弦退火一个完整周期optimizerSGD momentum0.9比 Adam 在微调任务上稳定initial lr0.01全模型微调常用范围weight_decay5e-4抑制过拟合T_max30与 epochs 一致eta_min1e-5学习率下限label smoothing0.1可选择配合 Mixup 有时能再提 0.5% 左右训练过程中需要关注验证损失曲线。因为 Mixup 混合了标签训练损失会比正常高 0.1 到 0.2 左右这不一定说明模型在变差重点看验证损失和准确率。如果验证准确率在第三个 epoch 后不再上升检查以下三个地方第一Normalize的均值和标准差是否与预训练模型匹配第二是否有类别样本数量极度不均衡导致模型偏向多数类可以打印每个类的召回率第三学习率是否过大导致损失震荡用CosineAnnealingLR时初始 lr 超过 0.01 往往会看到训练曲线陡升陡降。5. 预测阶段的两种写法与精度验证技巧训练完成后最终目标是拿到一张图片直接输出 12 类中的类别。预测写法常见的坑有两个一是忘记执行model.eval()导致 BatchNorm 和 Dropout 行为不一致二是预处理与训练不一致比如验证用了中心裁剪预测时却随机裁剪。下面两种写法都考虑到了这些点。5.1 单张图片路径预测def predict_single(model, image_path, class_names, devicecuda): model.eval() 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]) ]) img Image.open(image_path).convert(RGB) x transform(img).unsqueeze(0).to(device) with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim1) conf, idx torch.max(prob, dim1) return class_names[idx.item()], conf.item()unsqueeze(0)是在 batch 维度上增加一维因为模型要求输入为 (B, C, H, W)。torch.no_grad()在推理阶段必须加它可以让前向传播不构建计算图减少显存占用并加速。返回的conf是 softmax 置信度阈值为 0.5 时如果置信度过低可以输出未识别这在实际部署中很必要因为模型对不属于 12 类的图片也会强行给出一个类别。5.2 批量文件夹预测批量预测更关心吞吐量需要借助 DataLoader 复用 num_workers 多进程加载。此时不能用 predict_single 逐张循环那样会频繁切换 CPU/GPU效率很低。def predict_batch(model, image_dir, class_names, batch_size64, devicecuda): model.eval() 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]) ]) dataset SeedlingDataset(image_dir, transformtransform) loader torch.utils.data.DataLoader(dataset, batch_sizebatch_size, shuffleFalse, num_workers4) results [] with torch.no_grad(): for imgs, _ in loader: imgs imgs.to(device) logits model(imgs) prob torch.softmax(logits, dim1) conf, idx torch.max(prob, dim1) results.extend(zip(idx.cpu().numpy(), conf.cpu().numpy())) return results这里直接复用了自定义的SeedlingDataset但它会要求目录下每个子文件夹代表一个类别。如果只是把所有待预测图片放在一个文件夹里没有按类别分目录就需要临时写一个ImageFileDataset只读取文件列表并返回图像不关心标签。批量预测结束后每个元素是 (类别索引, 置信度)再通过class_names映射即可。5.3 用混淆矩阵验证预测结果单张预测和批量预测的准确率都依赖验证集评估方式。除了整体 accuracy 外我建议额外输出混淆矩阵确认 12 个类别里哪几类互相容易混淆。植物幼苗中有些类别在子叶阶段几乎长的一样模型混淆是正常的。实现时可以用 sklearn 的confusion_matrix也可以手动计算。手动方式更直观每行代表真实类别每列代表预测类别对角线占比越高说明该类别可分性越好。如果发现某两类交叉错误超过 10%不要急着加数据先检查真实标注是否准确其次再考虑针对这两个类别做 Oversampling 或生成额外合成样本。这个技巧比反复调学习率更能提升最终部署效果。本文还有配套的精品资源点击获取