简介MobileViT v3 的 PyTorch 实现代码包面向需要在移动设备上部署高效视觉模型的算法工程师与研究者。该实现以 MobileViT 架构为基础将局部卷积与全局 Transformer 自注意力结合在新版本中修正了官方代码加载预训练模型的兼容问题开发者拿到后可直接加载 ImageNet 权重进行迁移学习或微调能显著缩短复现与二次开发的时间。资源包共包含 607 个文件体积约 37.19MB主体为 264 个 py 源码与 221 个 pyc 编译文件同时带有 74 个 yaml 模型结构配置、23 个 md 说明文档以及 pt 预训练权重文件目录按训练、验证、部署等模块划分pyc 编译文件可帮助快速加载运行md 文档可用于查阅模型结构与使用说明。目前已有 1074 人学习/下载适合希望快速上手轻量级视觉 Transformer 或基于 MobileViT v3 做自定义分类任务的 PyTorch 用户。借助这份实现可先用官方预训练权重初始化模型再替换分类层适配新数据集配套的训练、验证与部署脚本配合 yaml 参数覆盖了从模型构建、权重加载到模型优化的主要流程可直接作为移动端视觉任务的起点。1. MobileViT v3 到底改了什么一个能跑 PyTorch 代码的轻量 Transformer 方案做移动端部署的人几乎都卡在同一个选择上CNN 模型轻、能上端但精度天花板明显Vision Transformer 精度高可参数和算力要求让手机根本跑不动。MobileViT v3 的思路是把两者缝起来在卷积骨架的若干层里插入全局建模模块既保留 CNN 的归纳偏置又拿到 Transformer 的长距离依赖能力。很多团队复现时只关心一件事有没有一份能直接跑通的 PyTorch 代码但代码只是结果真正决定你落地成不成的是 unfold/fold 的维度处理、注意力的摆放位置、以及融合层的设计。这篇笔记会从原理讲到可运行的实现把训练配置和踩坑记录一并给你。2. MobileViT v3 的设计拆解unfold/fold、注意力分组与通道分配2.1 先搞清楚 unfold/fold 在做什么它不是 ViT 的 patch embeddingViT 的做法是把图片切成不重叠的 patch每个 patch 拉平成长度为patch_size × patch_size × C的向量再过一个线性层变成 token。MobileViT 走的是另一条路它保留卷积特征图用 fold/unfold 把特征图“拆开再重组”。假设输入特征图是(B, C, H, W)patch_size 取 2。先把 H、W 分别切成 Hp、Wp 个格子每个格子内部是 2×2。Unfold 的意思是把每个格子里相同位置的像素抽出来放在同一个序列里。也就是说所有格子里的“左上角像素”组成序列 1“右上角像素”组成序列 2依此类推。序列长度不是格子里的像素数而是格子总数 Hp × Wp。这么做的好处是每个 token 的注意力范围覆盖整张特征图上同一相对位置的像素点相当于用很小的计算换来了全局视野。Fold 就是反向操作把处理完的序列按位置放回特征图。PyTorch 里没有现成的“MobileViT unfold”通常手动 reshape permute 来实现这也是新人最容易翻车的地方。2.2 v3 相比 v1/v2 的关键改动注意力分组、通道重分配和融合层v1 版本在每个 MobileViTBlock 里塞一个全局 Transformer计算量随分辨率增长明显。v2 用线性注意力替代 softmax 注意力进一步压低开销。v3 在工程复现里通常把精力放在三件事上第一注意力分组。把通道拆成若干组每组独立做自注意力避免单头 Attention 的 Q、K 矩阵维度过大。这相当于在计算量和表达力之间做了折中。第二通道重分配。v3 更强调用宽度倍率width multiplier控制每个 stage 的通道数浅层保持窄通道深层加宽。相比 v1 的均匀设计v3 的配置更像 MobileNetV3 的风格——短而宽的阶段放在网络后半段。第三融合层的升级。v1 是把局部分支和全局分支 concat 后过 1×1 卷积v3 常在融合后加残差连接甚至在部分实现里用 SE 模块对融合结果做通道注意力加权。残差连接能让梯度在局部和全局两条路径之间走得更稳训练初期不容易因为 Transformer 分支的加入而崩 loss。2.3 为什么浅层保持纯卷积深层才插入 Transformer这是 MobileViT 系列最核心的设计直觉。浅层特征是边缘、纹理局部相关性极强卷积 3×3 的窗口已经足够表达强行上 Transformer 除了增加计算量还会丢掉 CNN 的平移等变性在小数据集上容易过拟合。深层特征语义性强需要跨区域的信息交互这时 Transformer 的价值才体现出来。我一般按下面的配置组织 backbone阶段输出尺寸操作通道数StemH/42 个 stride2 的 3×3 卷积32→64Stage1H/42 个 MobileViTBlock96Stage2H/8下采样 4 个 MobileViTBlock192Stage3H/16下采样 3 个 MobileViTBlock384这个结构里Transformer 只出现在 Stage1 之后Stem 和第一个下采样层保持纯卷积。输入 256×256 的图Stage3 的特征图是 16×16patch_size 取 2 时序列长度只有 64注意力矩阵是 64×64几乎不占显存但语义交互能力已经覆盖了整图。3. 手写 MobileViT v3 的 PyTorch 代码核心 block 与 backbone 组装3.1 动手前先理清张量走线从 (B, C, H, W) 到 (B*P, N, C) 再到 fold 还原要实现 unfold先要确定中间张量的形状变化。设 B 为 batchC 为通道数HpH/phWpW/pwNHp×Wp 为格子总数Pph×pw 为每个格子内的像素个数。常见做法是先把特征图 reshape 成(B, C, Hp, ph, Wp, pw)再 permute 成(B, ph, pw, C, Hp, Wp)最后 reshape 成(B*P, C, N)转置后变成(B*P, N, C)。这里序列长度是 Nbatch 维被放大了 P 倍。把 P 当作 batch 的一部分意味着每个格子内的不同位置各自独立过 Transformer互不干扰这正是 MobileViT 与 ViT 的本质差异。3.2 实现多头自注意力与 MobileViTBlock 核心类下面这段代码是 block 的完整实现包含了局部卷积、unfold、多头注意力、fold、融合五部分import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadSelfAttention(nn.Module): 多头自注意力支持 key padding mask。 输入形状: (B, N, C)B 是 batch 与 P 合并后的维度。 def __init__(self, dim, heads4, dropout0.0): super().__init__() self.heads heads self.scale (dim // heads) ** -0.5 self.qkv nn.Linear(dim, dim * 3) self.attn_drop nn.Dropout(dropout) self.proj nn.Linear(dim, dim) def forward(self, x, maskNone): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.heads, C // self.heads) qkv qkv.permute(2, 0, 3, 1, 4) # (3, B, heads, N, dh) q, k, v qkv[0], qkv[1], qkv[2] attn (q k.transpose(-1, -2)) * self.scale if mask is not None: # mask 形状: (B, N)1 表示有效位置0 表示 padding attn attn.masked_fill(mask[:, None, None, :] 0, float(-inf)) attn attn.softmax(dim-1) if mask is not None: # 把 padding token 对应位置的行输出置零避免污染后续层 attn attn.masked_fill(mask[:, None, :, None] 0, 0.0) attn self.attn_drop(attn) x (attn v).transpose(1, 2).reshape(B, N, C) return self.proj(x) class MobileViTBlock(nn.Module): MobileViT v3 核心模块局部分支 全局 Transformer 分支 融合。 in_channels: 输入特征图通道数 d_model: Transformer 内部维度一般与 in_channels 相同 patch_size: (ph, pw)默认 2x2 def __init__(self, in_channels, d_model, patch_size(2, 2), heads4, dropout0.0): super().__init__() self.patch_size patch_size self.d_model d_model # 局部分支3x3 卷积聚合近邻上下文 self.local_conv nn.Sequential( nn.Conv2d(in_channels, d_model, 3, padding1, biasFalse), nn.BatchNorm2d(d_model), nn.SiLU(inplaceTrue), ) # 全局分支1x1 卷积把特征图投影到 d_model self.proj_conv nn.Conv2d(in_channels, d_model, 1, biasFalse) self.norm nn.LayerNorm(d_model) self.attn MultiHeadSelfAttention(d_model, headsheads, dropoutdropout) self.ffn nn.Sequential( nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model * 2, d_model), ) # 融合层concat 局部与全局后压缩回 d_model self.fusion nn.Sequential( nn.Conv2d(d_model * 2, d_model, 1, biasFalse), nn.BatchNorm2d(d_model), nn.SiLU(inplaceTrue), ) def forward(self, x): local_out self.local_conv(x) # 局部特征 h self.proj_conv(x) # 全局分支入口 B, C, H, W h.shape ph, pw self.patch_size # 1) pad 到 patch_size 的整数倍 Hp (H ph - 1) // ph Wp (W pw - 1) // pw H_pad, W_pad Hp * ph, Wp * pw pad_bottom H_pad - H pad_right W_pad - W if pad_bottom or pad_right: h F.pad(h, (0, pad_right, 0, pad_bottom)) mask torch.zeros(B, H_pad, W_pad, deviceh.device, dtypetorch.long) mask[:, :H, :W] 1 else: mask None # 2) unfoldreshape permute 成 (B*P, N, C) h h.view(B, C, Hp, ph, Wp, pw) h h.permute(0, 3, 5, 1, 2, 4) # (B, ph, pw, C, Hp, Wp) h h.reshape(B * ph * pw, C, Hp * Wp) h h.transpose(1, 2) # (B*P, N, C) # 3) Transformerattention FFN残差连接 shortcut h h self.norm(shortcut) h self.attn(h, maskmask) h h shortcut shortcut h h self.ffn(h) h h shortcut # 4) fold还原回特征图 h h.reshape(B, ph, pw, C, Hp, Wp) h h.permute(0, 3, 4, 1, 5, 2) # (B, C, Hp, ph, Wp, pw) h h.reshape(B, C, H_pad, W_pad) if pad_bottom or pad_right: h h[:, :, :H, :W] # 裁掉 padding # 5) 融合局部与全局 out self.fusion(torch.cat([local_out, h], dim1)) if out.shape[1] x.shape[1]: out out x # 输入输出通道一致时加残差 return out这里有几个要点需要说明。permute(0, 3, 5, 1, 2, 4)看起来复杂实际只做了一件事把 Hp、Wp 这两个“格子坐标”放在 C 前面把 ph、pw 这两个“格内偏移”移到 batch 维。这样 reshape 之后每个格内偏移就是一个独立的子 batch序列维度上只剩 NHp×Wp。mask的处理值得单独说。当输入尺寸不是 patch_size 的整数倍时右侧和下侧会出现 padding。如果不把 mask 传给注意力padding 区域的像素照样参与注意力加权模型会学着把权重放在没有意义的零值上轻则注意力图出现竖线重则精度掉 1 到 2 个点。FFN 里用了两层线性加 GELU中间维度扩到 d_model 的 2 倍这是轻量模型常用的比例。我见过有人把它扩到 4 倍效果提升不明显参数量和延迟却涨了不少移动端部署不建议这么做。3.3 组装 MobileViT v3 backboneStem、下采样与分类头核心 block 有了组装 backbone 就相对直接。按 2.3 的配置表我把 stem、三个 stage、分类头串起来class ConvNormAct(nn.Module): 卷积 BN SiLU 的组合MobileNet 系列的标准单元 def __init__(self, in_c, out_c, kernel_size3, stride1): super().__init__() self.conv nn.Conv2d(in_c, out_c, kernel_size, stridestride, paddingkernel_size // 2, biasFalse) self.bn nn.BatchNorm2d(out_c) self.act nn.SiLU(inplaceTrue) def forward(self, x): return self.act(self.bn(self.conv(x))) class DownSample(nn.Module): 下采样stride2 的 3x3 卷积通道数翻倍 def __init__(self, in_c, out_c): super().__init__() self.conv ConvNormAct(in_c, out_c, kernel_size3, stride2) def forward(self, x): return self.conv(x) class MobileViTv3(nn.Module): MobileViT v3 分类网络输入 3 通道图像输出 num_classes 分类 def __init__(self, num_classes1000, width_mult1.0, dropout0.1): super().__init__() m width_mult # 宽度倍率0.5/0.75/1.0 分别对应不同规格 self.stem nn.Sequential( ConvNormAct(3, int(32 * m), 3, stride2), ConvNormAct(int(32 * m), int(64 * m), 3, stride2), ) # stage 配置: (通道数, block 数量, d_model, patch_size) self.stage1 nn.Sequential( MobileViTBlock(int(64 * m), int(96 * m), patch_size(2, 2)), MobileViTBlock(int(96 * m), int(96 * m), patch_size(2, 2)), ) self.down2 DownSample(int(96 * m), int(192 * m)) self.stage2 nn.Sequential( MobileViTBlock(int(192 * m), int(192 * m), patch_size(2, 2)), MobileViTBlock(int(192 * m), int(192 * m), patch_size(2, 2)), MobileViTBlock(int(192 * m), int(192 * m), patch_size(2, 2)), MobileViTBlock(int(192 * m), int(192 * m), patch_size(2, 2)), ) self.down3 DownSample(int(192 * m), int(384 * m)) self.stage3 nn.Sequential( MobileViTBlock(int(384 * m), int(384 * m), patch_size(2, 2)), MobileViTBlock(int(384 * m), int(384 * m), patch_size(2, 2)), ) self.pool nn.AdaptiveAvgPool2d(1) self.head nn.Sequential( nn.Dropout(dropout), nn.Linear(int(384 * m), num_classes), ) # 常用初始化卷积用 kaiming线性层用截断正态 for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) elif isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std0.02) def forward(self, x): x self.stem(x) x self.stage1(x) x self.down2(x) x self.stage2(x) x self.down3(x) x self.stage3(x) x self.pool(x).flatten(1) x self.head(x) return xstage1 里第一个 block 的输入是 64 通道d_model 是 96输入输出通道不一致所以残差自动跳过第二个 block 输入输出都是 96残差生效。这个细节很多人没注意导致浅层梯度流通不畅。stage2 和 stage3 内部全部是同通道残差全程生效。初始化这里我用了混合策略卷积走 kaiming线性层走 trunc_normal。Transformer 分支如果不做特殊初始化在训练前期会输出方差偏大的结果直接叠加到卷积特征上容易让 loss 飞掉。trunc_normal 的 std0.02 已经足够保守。3.4 立即验证打印输出 shape、统计参数量与显存写完之后先别急着训练跑一段最小验证脚本确认每个 stage 的输出 shape 符合预期if __name__ __main__: model MobileViTv3(num_classes100, width_mult1.0) x torch.randn(2, 3, 256, 256) y model(x) print(输出:, y.shape) # 期望 (2, 100) total sum(p.numel() for p in model.parameters()) trainable sum(p.numel() for p in model.parameters() if p.requires_grad) print(f参数量: {total / 1e6:.2f}M, 可训练: {trainable / 1e6:.2f}M) # 估算显存占用 torch.cuda.empty_cache() if torch.cuda.is_available(): model model.cuda() x x.cuda() torch.cuda.reset_peak_memory_stats() with torch.no_grad(): y model(x) print(f峰值显存: {torch.cuda.max_memory_allocated() / 1024**2:.1f} MB)width_mult1.0 时参数量大致在 14M 到 16M 之间具体数值会随 PyTorch 版本和初始化方式有轻微浮动。256×256 输入、batch2 的峰值显存通常在 500MB 到 800MB 之间移动端 GPU 完全可以接受。如果显存吃紧把 width_mult 降到 0.75参数量会直接砍掉差不多一半精度损失通常在 1 个点以内。4. 训练配置与超参让 MobileViT v3 在 256×256 输入下稳定收敛4.1 数据增强的取舍移动端模型的泛化从增强来轻量模型参数少在小数据集上容易欠拟合而不是过拟合因此数据增强反而应该做得比大模型更重。224 到 256 的输入尺寸之间我建议选 256因为 MobileViT 的 Transformer 分支在低分辨率下计算量本来就不大大一点的分辨率能让全局建模优势发挥出来。我的常用增强组合是 RandomResizedCrop(256) 加 RandomHorizontalFlip 加 ColorJitter强度不要开满。强随机擦除Random Erasing在 MobileViT 上有时候会掉点因为 Transformer 分支对 token 缺失更敏感不像纯 CNN 那样扛得住。用不用 mixup 取决于数据集量级万级以下的数据集建议用百万级数据集可以不折腾。4.2 优化器、warmup 与余弦退火三个必调参数Transformer 分支和卷积分支共存时最忌讳用大学习率硬训。卷积对学习率不敏感但 Attention 的 QKV 投影层在训练初期梯度方差大学习率稍高 loss 就起飞。我一般直接上 AdamW把 betas 设成默认值 (0.9, 0.999)weight decay 给 0.05。warmup 是必须的不是可选。5 个 epoch 的线性 warmup 足够让 Transformer 分支的权重从初始化状态平滑过渡到正常更新区间。warmup 之后接余弦退火到初始学习率的 1/10 或 1/100比固定学习率多出明显收益。总训练 120 个 epochwarmup 5 个 epoch余弦周期剩下 115 个 epoch。import math import torch from torch.optim.lr_scheduler import LambdaLR EPOCHS 120 WARMUP_EPOCHS 5 BASE_LR 1e-3 def lr_lambda(epoch): if epoch WARMUP_EPOCHS: return (epoch 1) / WARMUP_EPOCHS progress (epoch - WARMUP_EPOCHS) / (EPOCHS - WARMUP_EPOCHS) # 余弦退火从 1 衰减到 0.01 return 0.01 0.99 * 0.5 * (1 math.cos(math.pi * progress)) optimizer torch.optim.AdamW(model.parameters(), lrBASE_LR, weight_decay0.05) scheduler LambdaLR(optimizer, lr_lambdalr_lambda)这里的lr_lambda是 LambdaLR 的调度函数输入当前 epoch输出学习率缩放系数。warmup 阶段系数从 1/WARMUP 线性涨到 1之后余弦从 1 降到 0.01。BASE_LR 是峰值学习率如果 batch 从 64 升到 256建议按比例调成 2e-3但不要超过这个值。4.3 训练循环骨架AMP 混合精度与 EMA 参数训练循环本身不复杂但有两个细节直接影响最终精度混合精度要用EMA 也要用。混合精度在 Transformer 分支上尤其值钱Attention 的矩阵乘是纯浮点运算FP16 能把显存占用砍掉一半训练速度提升 30% 以上。EMA 相当于给模型买了一颗后悔药。训练后期 loss 震荡时EMA 参数往往比当前参数更接近最优解推理时把 EMA 权重拷出来用通常能稳定提升 0.3 到 0.5 个点。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() def train_one_epoch(model, loader, optimizer, use_ampTrue): model.train() total_loss 0.0 for images, labels in loader: images images.cuda() labels labels.cuda() optimizer.zero_grad() if use_amp: with autocast(): outputs model(images) loss F.cross_entropy(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: outputs model(images) loss F.cross_entropy(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) return total_loss / len(loader.dataset)EMA 的实现建议单独维护一份权重副本列表每步训练后按ema_decay0.9999更新等训练完再加载回模型。如果训练中途想评估验证集直接同步一次 EMA 权重再评估评估完记得恢复原始权重否则 EMA 的统计意义会被打乱。4.4 完整超参参考表参数推荐值说明输入尺寸256×256不建议低于 224全局建模收益会下降Batch Size64与 BN 统计量有关见第 5 章峰值学习率1e-3AdamW warmup 配合使用Weight Decay0.05偏大的 wd 对 Transformer 分支更友好Warmup Epochs5线性 warmup总 Epochs120余弦退火末期降到 1e-5 量级EMA Decay0.9999训练完成后用 EMA 权重推理混合精度开启显存减半速度提升明显这套超参在 ImageNet 量级数据上训练 120 个 epoch精度基本能收敛到接近该规格模型的理论上限。如果数据是自定义的小数据集总 epoch 可以缩到 60warmup 缩到 3其余不用动。5. MobileViT v3 踩坑记录5 个复现翻车现场与修复5.1 unfold 维度对不上shape 漂移与 permute 顺序混乱现象forward 跑到 reshape 时报RuntimeError: shape [B, 3, 5, 1, 2, 4] is invalid for input of size ...或者输出 shape 比输入多一个维度怎么调都不对。原因把(B, C, Hp, ph, Wp, pw)permute 成(B, ph, pw, C, Hp, Wp)这一步目标维度的顺序必须和 reshape 前的 view 严格一致。很多人习惯直接照抄网上的 permute 参数但没注意到自己前面的 view 顺序是(B, C, Hp, ph, Wp, pw)而不是(B, C, Hp, Wp, ph, pw)。解决遇到 shape 问题先不要猜把每一行 reshape/permute 之前的 shape 都打印出来对照注释确认维度语义。最保险的验证办法是输入一个 1×1×4×4 的张量手算一遍期望输出代码跑通后再换真实尺寸。5.2 加了 Transformer 反而掉点全局分支的初始化与学习率现象同样的数据、同样的训练轮数跑了 MobileViTBlock 的模型精度还不如纯 CNN 的对照组loss 曲线看着也正常但 val acc 就是上不去。原因Transformer 分支的线性层如果沿用默认初始化通常是均匀分布在训练前期输出的方差偏大和局部卷积分支 concat 之后融合层的 BN 被迫适应一个剧烈波动的分布特征被“带偏”。另一个常见原因是学习率太大Attention 的梯度爆炸过一次之后权重就再也回不到良性区间。解决把线性层换成trunc_normal_(std0.02)初始化并把峰值学习率压到 1e-3 以下。如果还是掉点把融合层的 BN 换成epsilon1e-5降低数值敏感性通常能救回来。5.3 显存爆在注意力而不是卷积patch_size 与输入分辨率的关系现象训练时 GPU 显存报警报错堆栈指向MultiHeadSelfAttention的q k.transpose那一行。换成更小的模型也没用。原因很多人只关注通道数忽略了序列长度 N 对显存的影响。NHp×Wp输入 512×512、patch_size2 时N256×25665536注意力矩阵是 64K×64K一个 batch 的 attention 矩阵就要占掉几个 GB。卷积层的显存是线性增长注意力是平方增长这个差距在高分辨率下被急剧放大。解决训练阶段把输入分辨率控制在 256 以下patch_size 不要小于 2。推理阶段如果必须处理高分辨率输入先在低分辨率下跑通流程再针对高分辨率专门做切块或窗口化注意力不要妄想直接硬训。5.4 小 batch 下 BN 震荡换 GroupNorm 还是凑 batch现象batch size 设成 8 或者 4 之后训练 loss 正常下降但验证集精度波动非常大同一个模型不同验证批次之间能差 2 到 3 个点。原因MobileViTBlock 里局部卷积和融合层用了 BatchNormBN 的均值方差统计在小 batch 下噪声很大。Transformer 分支本身对 batch 不敏感但融合层把两个分布叠在一起BN 统计量一旦不准整体特征分布就跟着抖。解决如果显存允许优先把 batch 凑到 32 以上。显存实在不够把 fusion 和 local_conv 里的 BN 换成 GroupNormgroup 数设为 8效果远好过硬凑 batch。注意 stem 和下采样层可以保留 BN这两处没有 Transformer 分支参与BN 在小 batch 下问题不大。5.5 pad 区域的注意力权重污染mask 忘传导致特征图出现竖条现象训练时 loss 正常但可视化 attention 权重或特征图时能明显看到右侧和下侧有一条高亮的带状区域分类精度在特定类别上莫名偏低。原因输入分辨率不是 patch_size 整数倍时pad 出来的零值区域参与了注意力计算。Attention 会自动学会把部分权重放到这些零值上因为它们“好预测”。这些位置原本没有语义信息却被当成了有价值的 token特征图边缘出现条纹就是证据。解决在MultiHeadSelfAttention里加 mask把 padding 区域的 key 屏蔽掉。mask 的构造要在 unfold 之前完成因为 padding 发生在前向传播的早期后面 reshape 时 mask 也要跟着变形维度不匹配很容易漏。我通常会在 block 的 forward 开头就生成 mask然后一路带到 attention确保不会遗忘。6. 把 MobileViT v3 改造成检测 backbone多尺度输出与验证分类模型验证通过后真正的落地往往在检测或分割任务上。把 MobileViT v3 接到检测头关键在于让 backbone 输出多尺度特征图。分类网络只有一个最终输出检测需要的是从 Stage1、Stage2、Stage3 分别引出特征再接 FPN 或 PAN 融合。class MobileViTv3Backbone(nn.Module): 从分类模型中抽出多尺度特征供检测头使用 def __init__(self, width_mult1.0): super().__init__() # stem、stage1、down2、stage2、down3、stage3 的构造 # 与 MobileViTv3 完全一致这里省略重复代码 pass def forward(self, x): c2 self.stem(x) # 下采样4倍适合小目标 c3 self.stage1(c2) # 下采样4倍 c4 self.down2(c3) c4 self.stage2(c4) # 下采样8倍中目标 c5 self.down3(c4) c5 self.stage3(c5) # 下采样16倍大目标 return c3, c4, c5这里 c3、c4、c5 分别对应 4 倍、8 倍、16 倍下采样是检测头最常用的三档尺度。改造时要把分类模型里的 pool 和 head 去掉forward 直接返回特征图列表。验证方法很简单输入 640×640 的图打印三个输出的 shape确认分别是 (B, 96, 160, 160)、(B, 192, 80, 80)、(B, 384, 40, 40)。我习惯在改造完先跑一次 CPU 推理确保 shape 稳定后再接 FPN。接检测头之后会暴露一个新问题检测任务输入分辨率高、batch 小第 5 章说的显存和 BN 问题会一起找上门。我会把 stage2 和 stage3 里的 Block 换成 GroupNorm 版本并严格控制训练输入不超过 512。最后补一句经验每次改结构前先在 256×256 输入下跑通 forward 和 backwardshape 对了再上大分辨率能省下大量排错时间。希望帮到你。本文还有配套的精品资源点击获取