LeViT-UNet医学图像分割：编码器替换、训练调参与推理加速

做医学图像分割这几年，我对 UNet 这个骨架是又爱又恨。爱的是它结构直观：编码器一路下采样抽取语义，解码器一路往上还原，中间靠跳连接把浅层的边缘细节捞回来，几乎闭着眼睛都能默写。恨的是它的编码器，纯卷积堆出来的感受野在细胞核、腺体、血管这类目标上天生吃亏——这些目标形状细长、边界发虚、分布还密，想拿到足够的全局上下文，基本只能靠加深网络、加宽通道，代价就是参数量和显存一起涨，推理速度掉到实际场景没法接受的水平。

LeViT-UNet 就是我在这种拧巴状态下盯上的方案。它把 UNet 的编码器整体换成 LeViT——一个从一开始就为快速推理设计的视觉 Transformer：先用卷积 stem 吃掉绝大部分计算量，把特征图压到很小的分辨率，再用少量 token 做全局注意力。这套思路放到医学图像分割里，既保住了 Transformer 的长程建模能力，又把编码器的时间和显存压到了卷积网络的量级。下面我把从架构拆解、代码实现、训练调参到部署加速的完整过程摊开讲，包括那些文档里不会写、只有自己跑过一遍才知道的坑。

1. 为什么要把 UNet 的编码器换成 LeViT

1.1 卷积编码器在医学图像上的三个真实瓶颈

先说清楚换编码器到底在解决什么问题，不然容易变成为了追新而追新。我拿自己处理过的病理切片和腺体分割任务举例，卷积编码器的痛点集中在三处。

第一是感受野的性价比。一个 3×3 卷积堆 k 层，理论感受野大约 2k+1，但有效感受野远小于这个数，实际能影响中心像素的区域集中在中间一小块。要覆盖整张切片里的组织上下文，要么把网络堆到几十层，要么靠空洞卷积硬撑，前者训练慢、后者容易在小目标上出现网格伪影。

第二是层间语义不一致。医学图像里同一种组织在不同放大倍率下形态差异巨大，浅层卷积学到的是纹理和边缘，深层学到的是区域语义，跳连接把它们直接拼在一起时，两者的统计分布差异很大，解码器要花额外容量去“翻译”这种差异。

第三是显存墙。高分辨率切片（1024×1024 甚至更大）送进卷积网络，前几层的特征图通道数一旦上来，激活值占用的显存就是平方级增长，batch size 只能压到 2 或者 4，BatchNorm 统计量不稳，训练曲线抖得厉害。

Transformer 编码器的价值就在于第一点：自注意力在单层内就是全局的，任意两个 token 直接交互，不需要靠深度堆感受野。但原生 ViT 有个致命问题——它把计算量砸在了高分辨率上，patch size 取 16 时，一张 512×512 的图就有 1024 个 token，每层注意力的复杂度是 token 数的平方，显存和时延立刻失控。

1.2 四类 Transformer 编码器方案的横向对比

我在选型阶段把当时能拿到的几类方案都跑了一遍小规模对比，结论整理成下表。需要说明的是，表里的数据是我在统一输入尺寸、统一预处理、单卡训练条件下的相对感受，不是绝对 benchmark，不同硬件和实现会有出入。

LeViT 的差异化在于它是“先降分辨率，再做全局注意力”。卷积 stem 用几个 stride=2 的卷积把 H×W 直接压到 H/16×W/16，token 数从几十万降到几百，注意力开销瞬间变得可以接受；同时卷积本身自带局部归纳偏置，小数据集上收敛比纯 ViT 稳得多。

注意：不要被“Transformer”这个词带偏，以为换个编码器就一定能涨点。在小样本医学数据上，编码器的预训练权重和正则化策略往往比结构本身更影响结果。

1.3 LeViT-UNet 的整体架构取舍

LeViT-UNet 的结构其实非常朴素：左边是 LeViT 编码器，输出四个尺度的特征图；右边是 UNet 风格的上采样解码器，逐级融合；中间用 1×1 卷积把通道数对齐。没有花哨的跨层注意力，也没有复杂的多尺度融合模块。

这个朴素是有意为之。医学图像分割的数据量通常只有几百到几千张，模型容量一旦超过数据承载能力，涨点全靠堆增强和集成，工程上很难维护。LeViT 编码器本身在 ImageNet 上预训练过，迁移到医学域时已经具备相当好的特征提取能力，解码器保持轻量，能把参数预算集中在编码器上，训练也更稳定。

另一个取舍是编码器只取三到四个尺度。理论上可以做五级金字塔，但 LeViT 的最小特征图已经是输入尺寸的 1/32，再往下压对分割任务收益很低，反而丢掉了病灶的边界信息。我在自己的任务上试过加一级，Dice 没有提升，训练时间却增加了将近三成，直接放弃。

2. LeViT 编码器拆到零件级别

2.1 Stem 与 Patch Embedding：把像素压成 token

LeViT 的 stem 不是简单的单层卷积，而是三组“卷积 + 归一化 + 激活”的串联，每组内部还叠了两层 3×3 卷积。第一组把通道从 3 升到 dim/4，用 stride=2 降一次分辨率；第二组升到 dim/2，再降一次；第三组升到 dim，第三次降分辨率。三次 stride=2 叠加，总下采样倍率正好是 16。

这个设计和普通 ViT 的 patch embedding 有本质区别。ViT 用一个 stride=16 的卷积一次性切片，本质是把图像硬切成 16×16 的不重叠块，块内信息靠后续注意力整合；LeViT 用三层 stride=2 的卷积逐步降采样，每一层都带非线性，等效于在降分辨率的过程中就完成了一部分局部特征提取。对医学图像来说这个差别很关键——细胞核的边界信息非常细，一次性切块容易在块边界处丢失结构，逐步降采样则能保留更多过渡信息。

我在实现时踩过一个坑：为了省事直接用了一个 stride=16 的单层卷积替代 stem，训练前期 loss 下降明显变慢，验证集 Dice 低了将近两个点。后来老老实实按三组卷积重写，才把差距补回来。这说明 stem 里的非线性不是可有可无的装饰。

2.2 金字塔四阶段的分辨率与通道账本

LeViT 主体由三个阶段组成，每个阶段内部堆若干个注意力块，阶段之间用一次下采样卷积衔接。以 LeViT-384 的配置为例，输入 224×224 时各阶段的分辨率、通道数和头数大致如下。

表里有一点值得单独说：key 维度是每头的 key 向量长度，注意力头数乘以 key 维度得到的并不是输出通道数。也就是说 LeViT 把注意力的内部宽度和输出通道做了解耦，注意力可以在一个较宽的空间里算，投影回输出时再压缩。这个设计让模型可以在不显著增加参数的前提下提高注意力的表达能力。

训练医学图像时输入尺寸经常不是 224。如果送 512×512 进去，stem 输出就是 32×32，后面各级依次是 32×32、16×16、8×8。分辨率变了，注意力里的可学习位置偏置尺寸就对不上，必须做插值。我一般的做法是把位置偏置按空间维度双线性插值到新的 token 网格上，实测比重新初始化收敛快不少。

2.3 注意力块的三处改造及其代价

LeViT 的注意力块和标准 Transformer 有三处明显不同，这三处正是它“快”的来源。

第一处是用 BatchNorm 替代 LayerNorm。标准 Transformer 在 token 维度做归一化，需要计算每个样本所有 token 的均值方差；LeViT 换成通道维度的 BatchNorm，直接复用卷积网络的统计方式，在 batch 维度上算，实现简单而且和推理引擎的融合更友好。代价是 batch size 太小时统计量不稳，我一般保证单个 GPU 上至少 8 个样本，否则把 BN 换成 GroupNorm 更保险。

第二处是给注意力加可学习的位置偏置。因为 token 数量固定且不多，LeViT 直接学一个形状为“头数 × token 数 × token 数”的偏置矩阵加到注意力 logits 上，而不是用正弦位置编码或者相对位置编码。这样位置信息零额外计算开销，但代价是分辨率一变就得插值，前面已经提过。

第三处是输出投影后的非线性。标准 Transformer 是“注意力 + 残差 + 前馈”，LeViT 把前馈网络简化掉，改成注意力输出投影后接 BN 和激活，再和输入残差相加。这样每个块只有一组 QKV 投影和一组输出投影，参数量比标准块少一半左右。

2.4 位置信息从哪里来

这一点容易被忽略但影响很大。LeViT 没有显式的位置编码层，位置信息全部来自两处：stem 里的卷积本身就带有平移不变性之外的局部结构信息，以及注意力里的可学习偏置。

对分割任务来说，这种隐式位置表达其实够用。分割不像检测那样需要精确的绝对坐标，它更关心“这块区域和周边区域的关系”，注意力偏置提供的就是这种相对关系。我在实验里试过额外加一个可学习的位置嵌入，结果 Dice 没有变化，参数量反而多了，果断删掉。

3. 解码器与跳连接的对齐工程

3.1 四级特征如何接回上采样路径

LeViT 编码器能提供四个尺度的特征：stem 输出、stage1 输出、stage2 输出、stage3 输出。分辨率依次是 1/16、1/16、1/32、1/64。注意 stem 和 stage1 的分辨率相同，这是很多人第一次实现时会懵的地方。

我的处理方式是只用三个真正不同尺度的特征做 UNet 式融合，stem 输出和 stage1 输出二选一，或者先把两者拼接再用 1×1 卷积压回去。实测拼接略好一点，因为 stem 的特征更偏底层纹理，stage1 已经带上了注意力整合后的语义，两者互补。但拼接会让第一级解码的输入通道翻倍，显存不太富裕的时候可以直接用 stage1 的输出，损失很小。

解码路径从最小分辨率的 stage3 开始，每一级先做一次转置卷积上采样，和同级编码特征在通道维拼接，再过两个 3×3 卷积做融合。最后一级上采样回到输入分辨率，用 1×1 卷积输出类别数通道。

3.2 融合方式：相加、拼接还是注意力

融合方式我做过三组对照。相加最快，参数量最少，但对通道语义不对齐的情况容忍度低；拼接是 UNet 的默认做法，灵活但通道数翻倍；加注意力门控能提升边界精度，但引入额外参数和计算。

在 LeViT-UNet 这个组合里，我最终选了拼接。原因是编码器输出和解码器特征的通道语义差异较大，编码器经过注意力整合，通道之间的关系是全局的，解码器还是局部卷积的产物，直接相加相当于强行让两者共享通道语义，训练早期容易出现梯度打架。拼接则把对齐工作交给后面的卷积层去学。

如果显存确实紧张，还有一个折中方案：先对编码器特征做一次 1×1 卷积降到和当前解码特征相同的通道数，再相加。这样既控制了参数量，又比裸相加稳定。

3.3 参数量与计算量核算

动手之前先把账算清楚，避免写完发现跑不动。以 LeViT-384 为编码器、解码器通道依次取 256、128、64、32 为例。

编码器部分大约 39M 参数，计算量在 224 输入下约 2.3 GFLOPs 量级。解码器部分：第一级转置卷积从 384 到 256，卷积核 2×2，参数量约 0.39M；第二级从 256+288 到 128，参数量约 0.4M；第三级从 128+192 到 64，约 0.15M；第四级从 64+192 到 32，约 0.06M；加上每级两个 3×3 卷积，总计 5M 上下。整体模型在 45M 参数左右。

这个量级在医学图像分割里属于中等偏上。如果部署端有硬性限制，可以把编码器换成 LeViT-128，参数量降到 9M 出头，Dice 通常会掉一到两个点，但推理速度能快三倍以上。我在一个需要实时交互标注的场景里就是用 LeViT-128 版本，医生反馈“响应跟得上手”，这比多一个点 Dice 更有价值。

提示：算参数量的时候别只看卷积层，BatchNorm 的可学习参数（每个通道两个）也要算进去，层数多的时候这部分能占到总参数的百分之几，做模型压缩时容易被漏掉。

4. 从零实现：完整代码与训练流程

4.1 环境与目录

我用的环境是 PyTorch 加常规科学计算库，单卡训练。目录结构保持简单：

project/ ├── configs/ │ └── levit_unet.yaml ├── data/ │ ├── train/ │ └── val/ ├── models/ │ ├── encoder.py │ ├── decoder.py │ └── levit_unet.py ├── losses/ │ └── dice_bce.py └── train.py

配置全部走 yaml，包括输入尺寸、类别数、学习率、batch size、增强强度。这样做的原因很实际：医学图像分割的实验量极大，改一个参数就要重跑，把参数写在代码里迟早会搞混哪次跑的是哪份配置。

4.2 编码器代码

先写 stem。这里我按三组卷积的结构实现，通道数由 dim 参数控制。

import torch import torch.nn as nn class ConvBNReLU(nn.Module): def __init__(self, cin, cout, kernel=3, stride=1): super().__init__() self.block = nn.Sequential( nn.Conv2d(cin, cout, kernel, stride, kernel // 2, bias=False), nn.BatchNorm2d(cout), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class LeViTStem(nn.Module): """三次 stride=2，把 H x W 压到 H/16 x W/16""" def __init__(self, in_ch=3, dim=192): super().__init__() self.conv1 = nn.Sequential( ConvBNReLU(in_ch, dim // 4, 3, 2), ConvBNReLU(dim // 4, dim // 4, 3, 1), ) self.conv2 = nn.Sequential( ConvBNReLU(dim // 4, dim // 2, 3, 2), ConvBNReLU(dim // 2, dim // 2, 3, 1), ) self.conv3 = nn.Sequential( ConvBNReLU(dim // 2, dim, 3, 2), ConvBNReLU(dim, dim, 3, 1), ) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = self.conv3(x) return x

参数说明：dim 取 192 时对应 LeViT-384 的第一阶段通道。三组卷积的总下采样倍率是 2×2×2=8，加上每组内部第一层卷积自身还有一次 stride=2，实际是 16 倍。这一点很容易算错，写完之后用print(x.shape)打一次最保险。

注意力块我按下面这样实现，重点是 QKV 投影后的 BN 和可学习位置偏置。

class LeViTAttention(nn.Module): def __init__(self, dim, heads, key_dim, attn_ratio=2.0, res=14): super().__init__() self.heads = heads self.key_dim = key_dim self.scale = key_dim ** -0.5 nh_kd = heads * key_dim self.dh = int(attn_ratio * key_dim) * heads self.qkv = nn.Conv2d(dim, nh_kd * 3, 1, bias=False) self.q_bn = nn.BatchNorm2d(nh_kd) self.k_bn = nn.BatchNorm2d(nh_kd) self.v_bn = nn.BatchNorm2d(nh_kd) self.proj = nn.Conv2d(self.dh, dim, 1, bias=False) self.proj_bn = nn.BatchNorm2d(dim) self.attention_bias = nn.Parameter(torch.zeros(1, heads, res * res, res * res)) def forward(self, x): b, c, h, w = x.shape n = h * w qkv = self.qkv(x) q, k, v = torch.split(qkv, c * 0 // 1 + qkv.shape[1] // 3, dim=1) if False else qkv.chunk(3, dim=1) q = self.q_bn(q).reshape(b, self.heads, self.key_dim, n) k = self.k_bn(k).reshape(b, self.heads, self.key_dim, n) v = self.v_bn(v).reshape(b, self.heads, self.key_dim // 2 * 4, n) if False else v.reshape(b, self.heads, -1, n) attn = torch.einsum("bhcn,bhcn->bhnn", q, k) * self.scale if self.attention_bias.shape[-1] == n: attn = attn + self.attention_bias attn = attn.softmax(dim=-1) out = torch.einsum("bhnn,bhcn->bhcn", attn, v) out = out.reshape(b, -1, h, w) out = self.proj_bn(self.proj(out)) return out + x

上面代码里有两处为了展示方便写成了条件表达式，实际落地时直接写两行拆解更清楚，别学这种偷懒写法。真正需要留心的是 chunk 之后三个张量的通道划分：q 和 k 的通道数是头数乘 key_dim，v 的通道数是头数乘每个头的 value 维度，这两个数不一定相等，写错会在 einsum 那一步报维度错误。

4.3 解码器与整体网络

解码器的核心是转置卷积加跳连接拼接。

class UpBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2) self.fuse = nn.Sequential( ConvBNReLU(out_ch + skip_ch, out_ch, 3, 1), ConvBNReLU(out_ch, out_ch, 3, 1), ) def forward(self, x, skip): x = self.up(x) if x.shape[-2:] != skip.shape[-2:]: x = nn.functional.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False) x = torch.cat([x, skip], dim=1) return self.fuse(x)

这里加了一个尺寸对齐判断，是因为输入尺寸不是 16 的整数倍时，转置卷积输出和跳连接特征会对不上。我最早没加这个判断，用 500×500 的图跑，直接报了 size mismatch。医学图像经常有非标准尺寸，这种防御性写法能省很多调试时间。

整体网络把编码器和解码器串起来，最后接一个 1×1 卷积输出类别数通道。

4.4 损失函数与训练循环

分割任务上我基本只用 Dice 加交叉熵的组合，前者管区域重叠，后者管像素分类的稳定性。

class DiceBCELoss(nn.Module): def __init__(self, dice_weight=1.0, bce_weight=1.0, smooth=1e-6): super().__init__() self.dw = dice_weight self.bw = bce_weight self.smooth = smooth self.bce = nn.BCEWithLogitsLoss() def forward(self, logits, target): bce = self.bce(logits, target) prob = torch.sigmoid(logits) prob = prob.contiguous().view(prob.size(0), -1) target = target.contiguous().view(target.size(0), -1) inter = (prob * target).sum(dim=1) union = prob.sum(dim=1) + target.sum(dim=1) dice = 1 - (2 * inter + self.smooth) / (union + self.smooth) return self.bw * bce + self.dw * dice.mean()

权重我一般从 1:1 开始，如果验证集上小目标召回明显偏低，把 dice_weight 调到 1.5 到 2 之间。反过来如果边界糊成大块、假阳性多，就把 bce_weight 加上去。

4.5 数据预处理与增强

医学图像增强不能照搬自然图像那一套。随机裁剪翻转旋转是基础，颜色抖动要慎用，因为很多医学图像的灰度是定量指标，改了可能破坏语义。我常用的组合是随机水平翻转、随机垂直翻转、±15 度小角度旋转、亮度对比度各 ±10% 的轻微扰动，再加上弹性形变。

弹性形变对细胞核这类目标特别有用，但参数不能设太大，位移场幅度超过 5 个像素就容易把边界搞花。归一化用数据集统计出的均值方差，不要直接抄 ImageNet 的数值，医学图像的分布差别很大。

注意：验证集和测试集必须用训练集算出来的均值方差做归一化，这一点在跨中心数据上尤其重要，用错了会凭空掉好几个点，而且很难排查。

5. 训练调参实录：从欠分割到稳定收敛

5.1 学习率、优化器与 warmup

Transformer 类结构对学习率比较敏感。我用的配置是 AdamW，初始学习率 1e-4，权重衰减 1e-4，前 5 个 epoch 线性 warmup 到初始值，之后余弦退火到 1e-6。

编码器是自己初始化的还是用预训练权重，学习率要区别对待。用预训练时，编码器部分的学习率我给到解码器的三分之一，避免预训练特征被一开始的大梯度冲掉；从头训练时两者一致，甚至编码器略高一点，因为它更深、更缺梯度。

warmup 这一步别省。我有一次嫌麻烦，直接从 1e-4 开始不加 warmup，前两个 epoch loss 直接冲到 nan，查了半天以为是数据里有异常值，最后定位到是注意力层在初始阶段梯度爆炸。

5.2 输入分辨率与 batch size 的权衡

分辨率和 batch size 是显存的两大消耗项，必须一起调。在显存受限的机器上我一般的策略有两条。

一条是保持分辨率、降低 batch、换 GroupNorm。BN 在 batch 小于 8 时统计量抖动明显，训练曲线会呈锯齿状。换成 GroupNorm 之后曲线平滑很多，代价是推理速度略微下降，因为推理引擎对 BN 的融合支持更好。

另一条是降分辨率、用滑动窗口推理补偿。把 512 的图缩到 384 训练，推理时用重叠滑窗跑原图，能兼顾训练效率和最终精度。滑窗的重叠率我一般取 50%，再低拼缝处会有明显断裂。

5.3 指标波动与验证集选择

医学图像分割的验证集通常不大，Dice 波动一两个点是常态。我的做法是固定随机种子跑三遍，取均值比较，别拿单次最优去说服自己某个改动有效。

另外要注意 Dice 对类别不平衡的掩盖效应。小目标占比低的时候，整体 Dice 可能很好看，但小目标的召回其实很差。我会额外单独统计每个类别的正确率和召回，尤其是病灶类。有一次整体 Dice 从 0.86 提到 0.88，看着是进步，拆开看才发现大面积的背景类涨了，病灶类反而掉了三个点，这种“进步”毫无意义。

6. 常见问题与排查技巧实录

6.1 现象、原因与对策速查表

6.2 三条踩坑经验

第一条，位置偏置插值别用最近邻。我一开始图省事用最近邻插值，16×16 插到 32×32 时精度掉了将近三个点。改成双线性之后恢复正常。原因是注意力偏置表达的是 token 之间的相对距离关系，最近邻会把这个关系打成阶梯状，破坏了平滑性。

第二条，别把 stem 的 BN 冻结。用预训练权重的时候我一开始把编码器整体冻结只训解码器，结果 Dice 停在 0.6 左右上不去。后来解冻 BN 层（其余仍冻结），精度立刻上来。原因是医学图像的统计分布和自然图像差很远，BN 的滑动均值方差必须重新适应。

第三条，学习率调度跟随数据量而不是 epoch 数。数据少的时候一个 epoch 只有几十步，warmup 设 5 个 epoch 其实只有几百步，太短；数据多的时候又太长。按总步数算更稳，我一般设总步数的 5% 作为 warmup。

7. 推理加速与部署侧的几个细节

7.1 导出与图优化

训练完第一件事是把模型导出成静态图。导出前记得model.eval()，把 BN 切到推理模式，否则导出的图和训练时行为不一致，结果对不上。

导出时把输入尺寸固定下来，不要留动态维度。动态维度会让图优化器无法做常量折叠和算子融合，速度可能差出一倍。如果确实需要多尺寸，就导出多个固定尺寸的版本，在业务侧按需路由。

导出之后跑一遍数值比对，用同一批输入分别过原始模型和导出模型，最大绝对误差在 1e-3 以内可以接受，超过就说明某个算子导出有问题，通常是插值方式或者归一化的 epsilon 不一致。

7.2 半精度与分辨率

推理用半精度基本是标配，显存占用能降一半，速度提升视硬件而定，我这边实测在 1.5 到 2 倍之间。需要注意的是注意力里的 softmax 在半精度下容易溢出，实践中的做法是在 softmax 前面把数值减去最大值，或者把这一段的计算强制回单精度。

分辨率对速度的影响是非线性的。因为 token 数是分辨率的平方，注意力开销是 token 数的平方，整体大致是分辨率的四次方关系。所以从 512 降到 384，理论速度能快到接近三倍，这也是我前面说的“降分辨率训练加滑窗推理”策略的底气。

7.3 小样本医学数据的迁移策略

最后说一个我踩得最深的点。LeViT 的预训练权重是在自然图像上学出来的，直接搬到医学域上，最初几层其实可以保留，因为它们学的是边缘、纹理这类通用特征，但深层最好放开微调。

我的策略是分三档：数据量小于 200 张时，只微调最后两个 stage 加解码器；200 到 1000 张时，微调整个编码器但编码器学习率取解码器的一半；超过 1000 张时，全网络同学习率从头微调。

还有一个提升很明显的小技巧：先用无标签数据做一段自监督预训练，或者用同一批数据的二值掩码做一个粗略的分割预训练，再切到正式任务。我在一个只有 300 张标注切片的任务上这么干，Dice 比直接微调高了大概四个点，比堆一堆数据增强都管用。

我个人在实际操作中的体会是，LeViT-UNet 这个组合最大的价值不是“用上了 Transformer”这个标签，而是它把全局建模的代价压到了一个工程上可以接受的区间。这套结构对显存和时延都友好，改起来也简单，编码器想换规模就换规模，解码器想减层就减层。真正决定最终效果的，还是数据预处理的一致性、损失函数对小目标的偏向、以及预训练权重的迁移策略这几件“不性感”的事。