做图像分类项目的人多半经历过这个阶段:手头只有几千张自己标的数据,想上一台算力一般的机器跑训练,又不想在服务端堆一张大显卡。这种约束下,PyTorch 加 MobileNetV2 基本是最先被想到的组合——模型小、预训练权重现成、改两行代码就能跑起来。但"能跑起来"和"能跑出可用精度"之间隔着一整套工程细节:数据管道怎么搭、分类头怎么改、冻结策略怎么排、BatchNorm 什么时候该停、增强组合跟归一化有没有错配,这些地方任何一处出问题,验证集精度都会卡在一个尴尬的数字上不动。
这篇内容按我实际做项目的顺序展开:先把模型的能力边界讲清楚,再讲数据管道、分类头改造、两阶段训练、损失与优化器配置,最后落到精度排查和导出量化。适合手里有几百到几万张图、要在普通硬件上训一个能上线的分类模型的人看。代码都是可以直接抄下来跑的形态,参数取值我也会说明为什么这么给。
1. 先把MobileNetV2的能力边界摸清楚,再谈训练
很多人拿到这个模型就直奔models.mobilenet_v2(pretrained=True),然后开始调学习率,调了两天发现精度上不去,回头怀疑是学习率的问题。实际上大多数时候问题出在对模型结构的理解不够,导致后面的策略选择从一开始就是错的。
MobileNetV2 在 2018 年前后被提出,核心解决的是"在移动端算力预算内,如何让卷积网络的表示能力尽量不损失"。它的设计思路跟早年 VGG 那种一味堆通道、堆层数的路线完全不同,走的是另一条路:靠结构上的信息复用和通道数控制来压缩计算量。理解这一点,才能知道它适合什么任务、不适合什么任务。
1.1 倒残差结构和线性瓶颈到底省在了哪里
普通残差块(ResNet 那种)是"两头宽中间窄":先 1x1 把通道降下来,再 3x3 卷积,再 1x1 升回去。MobileNetV2 反过来了,是"两头窄中间宽":先用 1x1 把通道数扩张t倍(默认t=6),在扩张后的高维空间里做 3x3 深度可分离卷积,最后再用 1x1 压回低维。这就是所谓的Inverted Residual(倒残差)。
为什么要先扩张再卷积?因为深度可分离卷积的逐通道卷积部分本身表达力有限,如果直接在低维空间做,特征容易被"压扁"到恢复不回来。先升维再卷积,相当于给每个通道留出足够的活动空间,之后再用线性层压回去。
这里有个关键细节:压缩回去的那一层 1x1 卷积后面不接 ReLU。这就是Linear Bottleneck(线性瓶颈)。原因是 ReLU 会把负值全部截断为零,在低维空间做这种截断,信息损失是不可逆的;而在高维空间做截断,信息还能保留在其他通道里。所以 MobileNetV2 的残差块结构是1x1(升维) → ReLU6 → 3x3 DW → ReLU6 → 1x1(降维) → 无激活 → 与输入相加。
这个设计带来的直接后果是:模型中间层的特征维度普遍偏低(最后一层卷积输出只有 1280 维,很多中间块只有 24、32 这样的通道数),所以它对细粒度视觉任务的特征区分能力天生不如 ResNet50 这类模型。如果你的任务是区分十几种外观非常接近的工业零件、或者区分几十种相似鸟类,MobileNetV2 的天花板就在那里,再怎么调参也顶不上去多少,这时候该换骨干网络而不是继续磨参数。
同理,ReLU6(把激活值上限截断在 6)是为了配合后续的定点量化。如果你打算最终部署到端侧做 INT8 推理,这个设计是加分项;如果没有量化需求,它也不会带来明显损失。
1.2 什么数据规模配得上它,以及预训练权重该怎么用
MobileNetV2 在 ImageNet 上训练过 1000 类,特征提取能力已经有了不错的底子。实际项目里几乎没人从随机初始化开始训它,迁移学习是默认选项,区别只在于冻不冻结、冻多少。
我自己的经验划分是这样的:
| 自有数据量 | 推荐策略 | 预期验证集精度区间 |
|---|---|---|
| 每类 20 张以下(总量 <500) | 完全冻结 features,只训分类头,强增强 | 取决于类间差异,通常 70%-85% |
| 每类 50-200 张 | 先冻结训头,再解冻后 3-4 个 stage 微调 | 85%-93% |
| 每类 500 张以上 | 全量微调,小学习率,配合强增强 | 92%-97% |
| 十万张以上 | 可以考虑从头训或者换更大骨干 | 视任务而定 |
需要说明的是,这些区间是针对"图像内容与 ImageNet 自然图像分布接近"的场景。如果你做的是医学影像、遥感图像、工业缺陷检测这类与自然图像差异很大的域,预训练权重的收益会明显缩水,冻结策略要更激进地往后半部分走,甚至要考虑在目标域数据上做一段无监督的域适应预训练。
关于预训练权重的加载,PyTorch 新版接口已经改用 weights 枚举,老代码里的pretrained=True会在未来版本被移除:
import torch import torch.nn as nn from torchvision import models from torchvision.models import MobileNet_V2_Weights weights = MobileNet_V2_Weights.IMAGENET1K_V1 model = models.mobilenet_v2(weights=weights) # 配套的预处理参数,一定跟权重来自同一个枚举 preprocess = weights.transforms() print(preprocess)这里有个容易被忽略的点:预处理参数必须跟权重版本严格对应。上面这个枚举对应的归一化是mean=[0.485, 0.456, 0.406]、std=[0.229, 0.224, 0.225],输入尺寸 224。如果你手动写了另一套归一化参数(比如按自己数据集统计出来的均值方差),而模型权重是在 ImageNet 统计量下训出来的,特征分布就会整体偏移,精度会莫名其妙掉几个点。要改归一化,就得跟着微调,不能只改推理端。
2. 数据管道搭得对不对,决定了后面一半的调参都是白费
我见过太多精度上不去的案例,最后定位到问题都在数据这一层:训练集和验证集用了同一套增强、归一化参数跟权重不匹配、num_workers设太大导致 IO 打满、类别极不平衡但没做采样。模型本身没问题,数据管道在拖后腿。
2.1 目录结构与Dataset的最省事写法
图像分类最省事的组织方式是按类别分文件夹,一个文件夹一类:
dataset/ ├── train/ │ ├── cat/ │ │ ├── 0001.jpg │ │ └── ... │ └── dog/ │ └── ... └── val/ ├── cat/ └── dog/这种结构直接喂给ImageFolder就行,类别映射会自动按文件夹名排序生成:
from torchvision import datasets train_ds = datasets.ImageFolder("dataset/train", transform=train_tf) val_ds = datasets.ImageFolder("dataset/val", transform=val_tf) assert train_ds.class_to_idx == val_ds.class_to_idx, "类别映射不一致,检查目录结构" print(train_ds.class_to_idx, len(train_ds.classes))那个 assert 不是多余的。如果训练集里某一类只有一张图,验证集里没有这一类,ImageFolder生成的class_to_idx就会错位,训练时标签看起来正常,推理时全部张冠李戴。这种 bug 排查起来极其费时间,一行断言能省掉半天。
DataLoader 的配置也有讲究:
from torch.utils.data import DataLoader train_loader = DataLoader( train_ds, batch_size=64, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, persistent_workers=True, )drop_last=True在训练集里建议开着,因为最后一个 batch 如果只有一两张图,BatchNorm 统计出来的均值和方差会非常不稳定,对训练有负面影响。persistent_workers=True能避免每个 epoch 重启 worker 进程,在图片数量不多、epoch 很快的场景下能明显提速。num_workers不是越大越好,一般设成 CPU 核数的 1/4 到 1/2,设太大反而因为进程切换开销变慢。
2.2 训练集和验证集的transforms必须分开写
这是最基础但最常被写错的地方。训练集要增强,验证集绝对不能增强,只能做 resize、center crop 和归一化。
from torchvision import transforms IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] train_tf = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.6, 1.0), ratio=(0.75, 1.333)), transforms.RandomHorizontalFlip(p=0.5), transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.05), transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), transforms.RandomErasing(p=0.25, scale=(0.02, 0.15)), ]) val_tf = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), ])几点说明。RandomResizedCrop的scale=(0.6, 1.0)比默认的(0.08, 1.0)保守,默认值会把图像裁得只剩一个角,对于目标占据画面主体的数据集来说,裁得太狠会导致大量样本失去主体信息,反而拖慢收敛。数据集越小、目标越居中,scale下限应该越高。
RandomErasing(随机擦除)放在Normalize之后,因为它是对归一化后的张量做填充,填充值是 0,正好对应归一化后的"均值灰"。放错位置会得到奇怪的填充色块。
注意:
RandomHorizontalFlip不是万能的。文字识别、左右手区分、带方向性的工业零件检测这类任务,水平翻转会直接改变标签语义。上增强之前先想清楚这个变换在你的任务上是否保持标签不变。
2.3 类别不平衡时的采样器与损失加权
真实数据集里类别分布很少是均匀的。如果某一类占了 70%,模型只要全猜这一类就能拿到 70% 的准确率,训练日志看起来很美,实际完全不可用。
两种处理思路,我一般优先用采样器:
import numpy as np from torch.utils.data import WeightedRandomSampler targets = np.array(train_ds.targets) class_counts = np.bincount(targets) class_weights = 1.0 / class_counts sample_weights = class_weights[targets] sampler = WeightedRandomSampler( weights=torch.DoubleTensor(sample_weights), num_samples=len(sample_weights), replacement=True, ) train_loader = DataLoader(train_ds, batch_size=64, sampler=sampler, num_workers=4)注意用了sampler就不能再设shuffle=True,两者互斥,同时给会直接报错。
采样器的好处是每个 batch 内的类别分布更均衡,BatchNorm 的统计也更稳定;代价是少数类样本会被重复采样,容易过拟合。所以配合使用的时候,少数类的增强强度可以再调高一点,或者加一点标签平滑来缓解。
另一种做法是给CrossEntropyLoss传weight参数,但这个参数的量级对训练稳定性影响很大,权重给得太极端会让梯度爆炸。我一般只用采样器,实在不行再两者结合,并且把损失权重控制在[0.5, 3.0]这个范围内。
3. 改造分类头:从1000类走到你自己的类别数
模型结构看完、数据管道通了,接下来是改造输出层。这一步技术上很简单,但初始化方式和替换位置的细节会影响收敛速度。
3.1 定位classifier的最后一层
MobileNetV2 在 torchvision 里的实现,classifier是一个Sequential,结构是Dropout(0.2) → Linear(1280, 1000)。要改类别数,有两种写法:
# 写法一:直接替换整个 classifier in_features = model.classifier[1].in_features # 1280 model.classifier[1] = nn.Linear(in_features, num_classes) # 写法二:保留 Dropout,重建 Sequential model.classifier = nn.Sequential( nn.Dropout(p=0.2), nn.Linear(1280, num_classes), )两种写法等价,我更倾向第一种,因为它不依赖硬编码的 1280 和 Dropout 概率,换骨干网络时改动最小。
有一种常见错误是把model.classifier整个换成单个nn.Linear,看起来也能跑,但把 Dropout 也一起删掉了。在小数据集上微调,这 0.2 的 Dropout 对抑制过拟合是有实际作用的,删掉之后验证集精度往往会掉一两个点,而且训练损失下降得比验证损失快很多。
3.2 两种初始化策略与它们的收敛差异
新分类头默认是 PyTorch 的 kaiming 均匀初始化,直接开训能work,但在学习率给得偏大的时候,前几个 epoch 的 loss 会剧烈震荡。三种初始化方式我实测下来的差异:
| 初始化方式 | 收敛速度 | 前期稳定性 | 适用场景 |
|---|---|---|---|
| 默认初始化 | 中 | 一般 | 数据量充足,学习率调得比较小 |
正态分布std=0.01 | 快 | 好 | 冻结训练阶段,学习率较大时 |
| 全零初始化 | 极慢 | 极好但不推荐 | 仅用于分析实验 |
我通常在冻结阶段用std=0.01的正态初始化,这一阶段学习率会给到 1e-3 这个量级,小方差初始化能有效避免初始阶段梯度把预训练特征打乱:
nn.init.normal_(model.classifier[1].weight, std=0.01) nn.init.zeros_(model.classifier[1].bias)偏置初始化为零是标准做法,不需要改。
有个反直觉的现象值得提一句:分类头刚换掉的时候,模型输出是接近随机的,此时的 loss 大概在
ln(num_classes)附近。如果你第一轮 loss 就远高于这个值(比如 20 类任务,初始 loss 达到 8 以上),说明输入数据或者标签有问题,不要继续训下去,先回头检查。
4. 两阶段训练:先冻结再解冻的具体调度
这是整个流程里对最终精度影响最大的环节。一次性全量微调也能跑,但收敛慢、容易在早期把预训练学到的通用特征破坏掉。分两阶段走,前期稳定、后期精细,是小数据集微调的标准做法。
4.1 阶段一:只训分类头,让新层先"站稳"
冻结特征提取器,只训练新加的分类头。这个阶段的目的是让随机初始化的分类头先适应预训练特征的分布,避免一上来就用大梯度把整个网络带偏。
for param in model.features.parameters(): param.requires_grad = False optimizer = torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3, weight_decay=1e-4, )这里的filter(lambda p: p.requires_grad, ...)不是必须的,但把冻结的参数排除在优化器之外可以减少内存占用和优化器状态更新的开销,参数多的时候差别很明显。
阶段一通常跑 3-5 个 epoch 就够了,判据是验证集精度曲线开始变平。这一步的精度上限就是"线性探针"的水平——即用预训练特征直接做线性分类能达到的精度。它能给你一个很重要的参考:如果线性探针就能到 85%,说明特征质量不错,后续微调的空间在剩下的 5-10 个点;如果线性探针只有 60%,说明目标域跟 ImageNet 差异较大,微调时需要解冻更多层。
4.2 阶段二:解冻后半部分,差异化学习率
阶段一跑完,解冻。但不是全部解冻,也不是用同一个学习率。
MobileNetV2 一共 19 个 inverted residual 结构(在features里按 block 组织),浅层学的是边缘、纹理这类通用特征,深层学的才是语义相关的特征。域差异大的时候,深层需要更大调整;域差异小的时候,深层浅层都可以小幅度动。
# 找到 features 里的 block 边界 blocks = list(model.features.children()) print(len(blocks)) # 19(首层 conv 加上 18 个 inverted residual 序列) # 解冻后 4 个 block for b in blocks[-4:]: for p in b.parameters(): p.requires_grad = True # 差异化学习率 head_params = list(model.classifier.parameters()) backbone_params = [p for p in model.features.parameters() if p.requires_grad] optimizer = torch.optim.AdamW([ {"params": head_params, "lr": 5e-4}, {"params": backbone_params, "lr": 5e-5}, ], weight_decay=1e-4)头部的学习率是骨干的 10 倍,这个比例在实践中比较稳。骨干学习率给到 5e-5 这个量级,既能让预训练特征往目标域偏移,又不至于把已经学好的东西整个冲掉。
如果验证集精度在解冻后不升反降,通常有两个原因:一是骨干学习率给太大,二是解冻的层数太多。可以先把骨干学习率砍一半试试,还不行就只解冻最后两个 block。
4.3 BatchNorm到底该不该跟着动
这是冻结训练里最大的坑,没有之一。
requires_grad=False只是不更新 BN 层的weight和bias这两个可学习参数,BN 层的 running_mean 和 running_var 仍然会在 forward 时更新,因为它们是 buffer 不是 parameter,不受requires_grad控制。而后向传播时,只要model.train()开着,BN 就会用当前 batch 的统计量做归一化。
后果是:阶段一的时候,骨干的卷积权重是冻结的,但 BN 的统计量在被你自己的数据集(batch size 可能只有 32、64)不断改写。ImageNet 训练时用的 batch size 是几百上千,统计量非常准;用你的小 batch 改写之后,统计量噪声变大,特征分布发生偏移,会导致阶段一的精度明显低于预期。
处理方式有两种。第一种是在阶段一让整个骨干保持在 eval 模式:
def set_bn_eval(module): if isinstance(module, nn.BatchNorm2d): module.eval() model.train() model.features.apply(set_bn_eval) # 骨干的 BN 全部锁在 eval注意顺序:先model.train()整体切训练模式,再用apply把骨干里的 BN 单独切回 eval。如果顺序反了,前面的train()会把 BN 又切回来。
第二种是干脆把 BN 的 momentum 调小,比如从默认的 0.1 改成 0.01,让统计量更新得更平缓。我一般用第一种,干脆利落,不需要额外调参。
到了阶段二解冻之后,被解冻的 block 里的 BN 就应该恢复train模式了,这时候统计量跟着一起微调是合理的,因为卷积权重也在变。只锁住还没解冻部分的 BN。
一个可以快速验证的排查方法:训练结束后,把模型切到
eval()前后分别算一次验证集精度。如果两者差距超过 3 个点,基本可以确定是 BN 统计量的问题,要么是锁得不彻底,要么是训练后期的 batch size 太小。
5. 损失函数、优化器与学习率曲线的组合拳
模型和训练流程搭好之后,剩下能提精度的就是这几样配置。它们之间不是独立的,取值要互相配合。
5.1 label smoothing与weight decay的实操取值
交叉熵的标签平滑对小数据集微调几乎总是正收益。原理是把 one-hot 标签从[0, 0, ..., 1, ..., 0]软化到[ε/K, ..., 1-ε+ε/K, ..., ε/K],防止模型对单个样本过度自信。
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)0.1 是常用起点。类别数很少(比如 2-5 类)的时候可以降到 0.05,因为平滑强度相对于类别数是偏大的;类别数上百的时候可以给到 0.1-0.15。平滑给太大(比如 0.3 以上)会导致模型欠拟合,训练损失降不下去。
weight decay 方面,AdamW的默认值是 0.01,在微调场景下偏大。我一般给 1e-4 到 5e-4。这里有个值得注意的细节:weight decay 不应该加在 BatchNorm 的参数和所有 bias 上。这些参数本身数量很少,正则化它们没有意义,反而会影响归一化的稳定性。
def split_params(model): decay, no_decay = [], [] for name, p in model.named_parameters(): if not p.requires_grad: continue if p.ndim <= 1 or "bn" in name.lower() or "bias" in name: no_decay.append(p) else: decay.append(p) return decay, no_decay decay_p, no_decay_p = split_params(model) optimizer = torch.optim.AdamW([ {"params": decay_p, "weight_decay": 1e-4}, {"params": no_decay_p, "weight_decay": 0.0}, ], lr=5e-5)p.ndim <= 1这个判断能一次性覆盖所有 bias 和 BN 的 weight(BN 的 weight 是一维的),比按名字匹配更可靠。这个细节很多教程会跳过,但它在小数据集上的影响能有一两个点。
5.2 余弦退火还是OneCycle,看你的epoch预算
学习率调度器的选择取决于你能跑多少个 epoch。
如果总 epoch 数在 30 以上,CosineAnnealingLR是最稳的选择,配合 warmup 效果更好:
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR warmup = LinearLR(optimizer, start_factor=0.1, total_iters=3) cosine = CosineAnnealingLR(optimizer, T_max=total_epochs - 3, eta_min=1e-6) scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[3])前 3 个 epoch 从 10% 学习率线性升到设定值,避免一开始的大梯度冲击;之后余弦下降到1e-6。这个组合我用了很久,几乎不需要调。
如果只能跑 10-15 个 epoch,OneCycleLR收敛更快:
from torch.optim.lr_scheduler import OneCycleLR scheduler = OneCycleLR( optimizer, max_lr=5e-4, total_steps=total_epochs * len(train_loader), pct_start=0.3, )注意OneCycleLR是按 step 更新的,不是按 epoch,所以要在每个 batch 之后调scheduler.step()。这一点写错的人特别多,按 epoch 调会导致学习率曲线完全不对。
5.3 混合精度能省多少,什么时候不该开
AMP 在支持 Tensor Core 的显卡上能省 40% 左右的显存、提速 20%-30%,对 MobileNetV2 这种小模型来说收益没那么夸张(因为模型本身不是瓶颈,数据加载才是),但显存紧张的时候还是值得开。
scaler = torch.amp.GradScaler("cuda") for images, labels in train_loader: images, labels = images.cuda(non_blocking=True), labels.cuda(non_blocking=True) optimizer.zero_grad(set_to_none=True) with torch.amp.autocast("cuda", dtype=torch.float16): logits = model(images) loss = criterion(logits, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()set_to_none=True比默认的置零更省内存,在大模型上差异明显,小模型上也有收益。
什么时候不该开?一是显卡不支持(老架构开了反而慢),二是你发现 loss 出现 NaN 并且排除了其他原因。MobileNetV2 里有 ReLU6 这种截断操作,在某些边界情况下 fp16 溢出会出问题,虽然概率不高。如果训练不稳定,先关掉 AMP 排除掉这个变量。
6. 精度卡住不动时的排查链路
前面配置都给对了,验证集精度还是在一个不上不下的位置震荡,这时候要有条理地排查,别乱调参。我一般按这个顺序走。
6.1 先确认模型真的在学习,而不是在背标签
第一步永远是过拟合一个小样本。从训练集里挑 20 张图,关掉所有增强,训 100 个 epoch,看能不能到 100% 精度。
# 用 Subset 取出 20 个样本 from torch.utils.data import Subset small_idx = list(range(20)) small_ds = Subset(datasets.ImageFolder("dataset/train", transform=val_tf), small_idx)如果这个都做不到 100%,说明代码里有 bug——可能是标签映射错了、梯度没传、优化器没生效、或者 loss 计算方式有问题。这一步能排掉 80% 的"精度上不去"问题。
6.2 归一化错配与通道顺序
如果我前面强调过归一化参数要跟权重对应,这里再补两个更隐蔽的。
第一个是通道顺序。ImageFolder读进来的图片经过ToTensor()是 RGB 顺序,没问题。但如果你自己写 Dataset 用 OpenCV 读图,默认是 BGR,忘了转的话模型看到的颜色是错的,精度会掉得莫名其妙。转换就一行cv2.cvtColor(img, cv2.COLOR_BGR2RGB),但忘了就很难查。
第二个是图片位深。有些相机导出的图是 16 位 PNG,ToTensor()会把它当成[0, 1]范围的浮点处理,结果整张图几乎全黑。可以用PIL.Image.open(path).mode检查一下,如果是I;16就得先做位深转换。
6.3 训练模式切换的遗漏
model.train()和model.eval()的切换,验证阶段忘了切eval()是最常见的低级错误。忘了切的话,验证时 Dropout 还在随机丢弃、BN 还在用 batch 统计量,精度会比实际低一大截,而且同一个模型两次验证结果不一样。
我习惯把训练和验证封装成两个函数,在函数内部第一行就做模式切换,这样不会漏:
def train_one_epoch(model, loader, criterion, optimizer, scaler, device): model.train() model.features.apply(set_bn_eval) # 阶段一专用,阶段二删掉这行 ... def evaluate(model, loader, criterion, device): model.eval() with torch.no_grad(): ...验证函数里的torch.no_grad()同样重要,不加的话显存会持续增长,跑到后面容易 OOM。
6.4 区分过拟合、欠拟合和"模型容量不够"
看训练损失和验证损失的曲线关系来定位:
| 现象 | 判断 | 处理方向 |
|---|---|---|
| 训练损失高,验证损失也高,且接近 | 欠拟合 | 解冻更多层、加大学习率、减少正则 |
| 训练损失持续下降,验证损失先降后升 | 过拟合 | 加增强、加 Dropout、早停、减 epoch |
| 训练损失降到很低,验证损失降不下去 | 过拟合或数据泄漏 | 检查验证集是否与训练集有重复图片 |
| 训练损失和验证损失都在低位平缓,验证精度仍不高 | 模型容量不足或标签噪声 | 换骨干网络、清洗标签 |
| 两者都剧烈震荡 | 学习率过大或 batch 太小 | 降学习率、增大 batch |
第三行那个"数据泄漏"值得单独说。很多人在网上爬数据或者从公开数据集里切分,训练集和验证集里出现了同一张图(或者同一张图的裁剪、翻转版本),验证精度会虚高,上线后一测就打回原形。做去重的时候,用感知哈希(pHash)比精确哈希更靠谱,能抓到缩放和压缩后近似重复的图。
还有一种情况是标注本身有问题。几百张图的任务里,如果标注错误率在 5% 以上,模型精度天花板基本就被锁死了。我做过一个实验:故意在 20% 的标签上制造错误,模型最高只能到 80% 左右,跟理论值完全吻合。所以精度上不去的时候,抽 50 张验证集样本人工核对一遍标签,往往比调一天参有用。
7. 导出与量化:把模型从notebook搬到产线
训练完拿到一个满意的精度,最后一步是导出。这一步的坑集中在"训练时的行为和导出后的行为不一致"。
7.1 ONNX导出的两个必填参数
import torch model.eval() dummy = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy, "mobilenetv2.onnx", input_names=["input"], output_names=["logits"], opset_version=12, dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}}, )三点注意。第一,导出前必须model.eval(),否则 BN 和 Dropout 的行为会被固化到计算图里,导出的模型和训练时不是同一个东西。第二,dynamic_axes要显式声明,不声明的话 batch 维度被固化成 1,线上多batch推理会直接报错。第三,opset_version至少 12,低于这个版本对某些算子支持不全,虽然 MobileNetV2 结构简单基本不会踩到,但保险起见给 12 或 13。
导出后一定要用 onnxruntime 跑一遍,跟 PyTorch 的输出对比:
import onnxruntime as ort import numpy as np sess = ort.InferenceSession("mobilenetv2.onnx") out_onnx = sess.run(None, {"input": dummy.numpy()})[0] out_torch = model(dummy).detach().numpy() print(np.abs(out_onnx - out_torch).max()) # 应该在 1e-4 以内如果差得很多,八成是漏了eval()或者输入预处理不一致。这个比对步骤不要跳过,我见过导出的模型因为 Dropout 没关,线上推理结果完全随机的情况。
7.2 TorchScript:拟合BN之后的图更干净
如果目标环境是 C++ 或者移动端,TorchScript 是更直接的选择:
scripted = torch.jit.script(model) scripted.save("mobilenetv2.pt")torch.jit.script和torch.jit.trace的区别值得说清楚:trace是跑一遍前向、把实际执行的操作记录下来,遇到 if/for 这种控制流会直接固化;script是真正解析代码生成 IR,能保留控制流。MobileNetV2 是纯前向结构,两者都能用,但script更安全。
导出 TorchScript 之后可以用torch.jit.freeze把 BN 的参数折叠进卷积,推理速度能再快一点:
frozen = torch.jit.freeze(scripted.eval())折叠之后模型里的 BN 层消失了,卷积的权重和偏置已经被改写。这一步对于做量化特别重要,因为 BN 在量化的时候处理起来麻烦。
7.3 静态量化与掉点补偿的实际做法
端侧部署想跑得更快,通常会做 INT8 静态量化。PyTorch 的流程是"准备 → 校准 → 转换"三步:
import torch.quantization as tq model.eval() model.qconfig = tq.get_default_qconfig("fbgemm") # x86 CPU 用 fbgemm,ARM 用 qnnpack tq.prepare(model, inplace=True) # 用校准集跑一遍,收集激活值的分布 with torch.no_grad(): for images, _ in calib_loader: # 100-500 张图足够 model(images) tq.convert(model, inplace=True)校准集的选择很关键:必须用真实分布的数据,而且要覆盖所有类别。如果校准集里某一类一张都没有,这一类的激活范围就没被观测到,量化后这一类的精度会掉得特别厉害。我一般从验证集里每类随机抽 5-10 张,凑够 200-500 张做校准。
量化后的精度掉点,正常情况在 0.5-2 个点之间。如果掉超过 5 个点,通常是三个原因:校准集太少或者分布不匹配、量化配置选错(x86 用 fbgemm,ARM 用 qnnpack,选反了性能反而下降)、模型里有不支持的算子被打断成浮点计算。
掉点实在压不下去的时候,可以上 QAT(量化感知训练),在训练时就插入伪量化节点,让模型提前适应量化误差:
model.train() model.qconfig = tq.get_default_qat_qconfig("fbgemm") tq.prepare_qat(model, inplace=True) # 正常训练 5-10 个 epoch,学习率给正常微调的 1/10 # ... model.eval() tq.convert(model, inplace=True)QAT 能把掉点压到 0.5 个点以内,代价是要重新训一轮。数据量小、掉点又不可接受的时候,这个成本是值得的。
提示:量化只对 CPU 和特定加速器有意义。在 GPU 上跑 INT8 的收益取决于硬件和推理框架的支持程度,不要想当然地认为量化在 GPU 上一定更快。做之前先跑个 benchmark。
最后分享一个我踩过的坑。早期做量化的时候,我在convert之后直接拿原模型和新模型对比输出 logits 的差异,发现差了 0.05 左右,以为出了问题。实际上量化后的 logits 数值本身就不可直接比较,正确的验证方式是比对预测类别的一致率,在验证集上跑一遍,看量化前后的预测结果重合度。只要重合度在 98% 以上、精度掉点在 1-2 个点内,就是正常的量化损失,不需要返工。
这个组合我从早期的分类小项目一直用到最近的端侧方案,中间换过几次骨干网络,最后又回到 MobileNetV2,原因很简单:在数据量不大、部署资源有限的场景里,它的性价比是最高的,而且生态成熟,从训练到量化的每一环都有现成工具。真正花时间的从来不是模型本身,而是数据管道、BN 统计量、归一化一致性这些琐碎的地方,出问题的九成也都在这里。