做个医学图像分割骨架网络天天在变从早年的FCN到U-Net再到各种Transformer变体但很多新手一上来就抱着U-Net论文死磕模型改了几天换到真实数据上效果还是不行。U-Net当然经典可它的编码器说到底还是纯卷积堆叠网络一旦加深梯度回传和特征复用都会出问题。这绝不是危言耸听我在自己的项目里试过把U-Net编码器从3层加到5层训练loss直接降不下去。后来把ResNet思路融进去改用ResUnet问题一下子清晰了——残差连接不仅让梯度传得更稳特征表达也变得更充分。这篇文章就用PyTorch从头搭一个ResUnet专门用来做医学图像分割环境配置、网络结构、数据处理、训练推理、踩坑记录全部写清楚代码直接可以抄。1. 为什么我放弃“死磕”U-Net转向ResUnet1.1 U-Net的痛点在哪里U-Net最核心的设计是编码器-解码器结构加上跳跃连接。编码器通过逐层下采样提取高层语义信息解码器逐步恢复空间分辨率跳跃连接把编码器各层特征直接拼到解码器对应层弥补下采样丢失的细节。这个思路在二维医学图像分割上确实很有效尤其器官、病灶这类边界相对清晰的目标。可问题在于原始U-Net的编码器基本就是“两次卷积 ReLU 池化”的重复堆叠。网络深度有限一旦尝试加深去提取更抽象的特征梯度在多层卷积之间反复相乘很容易出现梯度消失。更麻烦的是深层卷积层如果学不到有效更新前面浅层学到的特征也会被后续层破坏这就是典型的退化问题。我在处理肝脏CT序列时用标准U-Net训练到30轮之后loss几乎不动加了残差结构后同一套数据和超参数loss还能继续稳定下降这个对比非常直观。另外医学图像本身有几个特点标注数据稀缺、类间不平衡严重、目标边界模糊。比如肿瘤分割中肿瘤区域往往只占整张图像的百分之几如果网络表达能力不够很容易把所有像素都判成背景。ResUnet通过残差连接增强特征复用能力让网络在有限数据下更容易学到判别性特征这也是我从U-Net转向ResUnet的根本原因。1.2 ResUnet的核心思路残差传给谁ResUnet并不是一个非常新的网络但它把两个成熟思路融合得很好一个是ResNet提出的残差学习另一个是U-Net的编码器-解码器结构。残差连接可以做得很细不只在编码器和解码器的每个Block内部加shortcut跳跃连接本身也可以理解为一种残差路径把浅层细节直接“传”给深层解码器。用公式来说普通卷积Block学习的是从输入x到输出y的直接映射而残差Block学习的是残差映射F(x)最终输出为y F(x) x。这样做的价值在于如果某个卷积层对特征的改变没有帮助网络可以直接把权重推向零让输出约等于输入而不会破坏已有特征。这一性质对医学图像分割特别重要因为病灶区域形状差异大、纹理细节丰富网络需要保留更多原始信息。我在实现ResUnet时参考了常见的设计编码器每个stage使用两个卷积层组成一个残差块卷积核大小保持3x3stride1padding1保证特征图尺寸不变下采样使用stride2的卷积或最大池化解码器每个stage先上采样再与对应编码器特征在通道维度拼接然后经过残差块。整体结构看起来就是“U-Net的骨架 ResNet的肌肉”。2. PyTorch环境搭建与工程准备2.1 用Anaconda创建独立环境动手写代码前先把PyTorch环境收拾干净。我强烈建议用Anaconda管理Python环境不要一股脑把包装到系统Python里。医学图像分割项目依赖版本往往互相牵扯今天装这个库、明天升级那个库没有独立环境迟早会出事。打开终端或者Anaconda Prompt执行以下命令创建Python 3.9环境conda create -n resunet python3.9 conda activate resunet为什么推荐Python 3.9因为PyTorch、MONAI、OpenCV、NumPy这些库对3.8到3.10的兼容性都很好3.11以后个别依赖可能会出现编译问题新手没必要在这种地方折腾。创建完环境后顺手把基础库装上pip install numpy opencv-python pillow matplotlib tqdm如果是做医学图像处理后面还需要用到SimpleITK或者MONAI但第一步先别急着全装等用到哪个装哪个避免环境里出现一堆用不上的包。2.2 安装PyTorchGPU版PyTorch安装这个环节看起来简单实际上很多人卡在这里。最稳妥的办法是访问PyTorch官网的Get Started页面根据操作系统和CUDA版本选择命令。这里要提醒一句首先要确认自己电脑的显卡支持CUDA用NVIDIA显卡最省心在命令行执行nvidia-smi就能看到驱动版本和支持的CUDA版本。如果是纯CPU环境学习也可以用CPU版但训练医学图像分割模型会非常慢建议至少用一个老款GTX 1660级别的卡。假设CUDA版本是11.8那么安装命令大致是pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118安装完成后务必验证一下这一步别偷懒import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果cuda.is_available()输出True说明GPU环境正常。很多人的问题出在明明装了CUDA版PyTorch却显示False这时候先检查NVIDIA驱动版本再检查conda环境是否激活最后检查是否装成了CPU版本。我在Win10上用Anaconda加PyCharm做这个项目时就碰到过PyCharm解释器指向了别的环境导致GPU不可用排查了半天才发现是解释器选错了。3. ResUnet网络架构逐层拆解3.1 从Residual Block开始ResUnet的基本单元是残差块我习惯先写一个小类保证编码器和解码器都能复用。代码不长但细节不少。import torch import torch.nn as nn class ResidualBlock(nn.Module): def __init__(self, in_channels, out_channels, stride1): super(ResidualBlock, self).__init__() self.conv1 nn.Conv2d(in_channels, out_channels, kernel_size3, stridestride, padding1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(out_channels, out_channels, kernel_size3, stride1, padding1, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) self.shortcut nn.Sequential() if stride ! 1 or in_channels ! out_channels: self.shortcut nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(out_channels) ) def forward(self, x): identity self.shortcut(x) out self.conv1(x) out self.bn1(out) out self.relu(out) out self.conv2(out) out self.bn2(out) out identity out self.relu(out) return out这里面最关键的是shortcut的设计。当输入和输出的通道数不一致或者特征图尺寸因为stride2被压缩时shortcut必须用1x1卷积调整维度和尺寸否则两个张量无法相加。我一开始图省事直接在shortcut里用nn.MaxPool2d处理下采样后来发现1x1卷积加stride的写法更灵活既能改变通道数也能下采样。3.2 编码器、解码器、跳跃连接完整网络结构可以拆成三个部分编码器、解码器、跳跃连接。编码器负责逐层提取语义信息我用四个stage每层通道数分别是32、64、128、256。这里要说明一下很多人一开始就把通道数设置成64、128、256、512照着ImageNet分类那套来实际在医学图像分割上不一定好因为医学数据量小参数量太大反而容易过拟合。我自己在多个数据集上试下来从32开始逐步翻倍效果和计算量的平衡最好。class ResUNet(nn.Module): def __init__(self, in_channels1, num_classes1): super(ResUNet, self).__init__() self.encoder1 ResidualBlock(in_channels, 32) self.encoder2 ResidualBlock(32, 64, stride2) self.encoder3 ResidualBlock(64, 128, stride2) self.encoder4 ResidualBlock(128, 256, stride2) self.bridge ResidualBlock(256, 512, stride2) self.decoder4 self._make_decoder(512, 256) self.decoder3 self._make_decoder(256, 128) self.decoder2 self._make_decoder(128, 64) self.decoder1 self._make_decoder(64, 32) self.final_conv nn.Conv2d(32, num_classes, kernel_size1) def _make_decoder(self, in_channels, out_channels): return nn.Sequential( nn.ConvTranspose2d(in_channels, out_channels, kernel_size2, stride2), ResidualBlock(out_channels * 2, out_channels) )等等这里有个容易混淆的地方。_make_decoder里我先用转置卷积上采样把通道数从in_channels变成out_channels但之后要和编码器跳跃连接拼接拼接后通道数变成out_channels * 2所以ResidualBlock的输入要设置成out_channels * 2。实际forward里必须先把编码器特征和上采样特征在通道维度拼接再送入残差块。def forward(self, x): e1 self.encoder1(x) # 32 e2 self.encoder2(e1) # 64 e3 self.encoder3(e2) # 128 e4 self.encoder4(e3) # 256 bridge self.bridge(e4) # 512 d4 self.decoder4[0](bridge) d4 torch.cat([d4, e4], dim1) d4 self.decoder4[1](d4) d3 self.decoder3[0](d4) d3 torch.cat([d3, e3], dim1) d3 self.decoder3[1](d3) d2 self.decoder2[0](d3) d2 torch.cat([d2, e2], dim1) d2 self.decoder2[1](d2) d1 self.decoder1[0](d2) d1 torch.cat([d1, e1], dim1) d1 self.decoder1[1](d1) out self.final_conv(d1) return out注意forward里的拼接顺序通常把编码器特征放在前面上采样特征放在后面或者反过来都行但一定要和训练/推理时的数据流保持一致。另外self.decoder4[0]这种写法依赖torch.cat时维度对齐如果输入图像尺寸不是16的倍数转置卷积之后的特征图和编码器特征图尺寸会不一致拼接就会报错。所以数据预处理阶段最好是统一resize到64的倍数比如256x256或者512x512。3.3 损失函数与评估指标为什么这样选医学图像分割里最常用的损失函数是Dice Loss它直接优化Dice系数比逐像素交叉熵更能应对类别不平衡。标准实现如下class DiceLoss(nn.Module): def __init__(self, smooth1e-6): super(DiceLoss, self).__init__() self.smooth smooth def forward(self, pred, target): pred torch.sigmoid(pred) pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() dice (2.0 * intersection self.smooth) / (pred.sum() target.sum() self.smooth) return 1 - dice这里的target是二值掩膜取值为0和1pred经过sigmoid后变成0到1之间的概率值。为什么不直接用交叉熵因为医学图像里前景占比经常低于10%交叉熵的梯度会被背景像素主导网络倾向于把所有像素预测成背景。Dice Loss把前景和背景看作一个整体来度量重叠程度对类别不平衡更鲁棒。评估指标我也习惯额外计算IoU和Dice系数方便和已有论文结果对比。训练时每个epoch结束后在验证集上计算这些指标保存Dice最高的一次权重而不是最后一轮权重这个细节后面还会提到。4. 医学图像分割实战从数据到可视化4.1 数据预处理不要盲目resize医学图像格式五花八门有的是二维PNG切片有的是三维NIfTI文件。这里先用最常见的二维切片场景举例。很多新手拿到数据直接cv2.resize到256x256然后丢进网络结果效果很差。原因在于盲目resize会破坏像素间距比例导致器官或病灶的形状变形。经典的医学图像处理工具SimpleITK或MONAI可以读取NIfTI里的spacing信息按照目标尺寸重采样但二维切片场景下如果标注本身已经是PNG那resize影响倒不大重点要保证所有样本使用相同尺寸。我自己常用的预处理流程是读取图像和对应掩膜统一转成灰度图。对图像做窗宽窗位调整比如CT图像通常关注[-200, 200] HU范围把范围外的值截断。归一化到[0, 1]公式是(img - min) / (max - min)。使用数据增强随机旋转、水平翻转、弹性形变等增强泛化性。如果内存充足可以读取所有数据放到NumPy数组里否则用torch.utils.data.Dataset按需读取。这里有一个关键点数据增强时图像和掩膜必须使用完全相同的变换。用torchvision.transforms里的RandomHorizontalFlip时传入的image和mask是分开处理的很容易出现图像翻转了、mask没翻转的情况导致训练标签错位。我建议用Albumentations库它专门支持同时变换image和mask大幅减少这类低级错误。import albumentations as A train_transform A.Compose([ A.RandomRotate90(), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.1), A.RandomBrightnessContrast(p0.2), ]) def preprocess(image, mask, size256): image cv2.resize(image, (size, size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (size, size), interpolationcv2.INTER_NEAREST) image image.astype(np.float32) / 255.0 mask mask.astype(np.float32) / 255.0 return image, mask注意mask的resize插值方式一定要用INTER_NEAREST如果用线性插值会引入中间灰度值破坏标签的类别语义。4.2 训练流程与关键参数训练一个分割模型参数选择比网络结构本身更影响最终效果。我一般在单卡GPU上使用batch size为8到16图像尺寸256x256编码器初始通道32总参数量约2000万显存占用不高。先把训练集和验证集切出来通常按8:2划分划分时按病人或样本ID分而不是按切片直接随机分避免同一个病人的相邻切片同时出现在训练集和验证集里导致数据泄露。这个细节对医学图像特别重要否则验证集的Dice会虚高真实场景泛化性能很差。优化器我推荐AdamW初始学习率1e-4并配合CosineAnnealingLR或ReduceLROnPlateau。早年我习惯用Adam后来发现AdamW的权重衰减处理更规范配合warmup效果更稳。训练代码如下import torch.optim as optim model ResUNet(in_channels1, num_classes1) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion DiceLoss() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100, eta_min1e-6) best_dice 0.0 num_epochs 100 for epoch in range(num_epochs): model.train() train_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device).unsqueeze(1) # (B,1,H,W) outputs model(images) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() * images.size(0) scheduler.step() # 验证 model.eval() val_dice 0.0 with torch.no_grad(): for images, masks in val_loader: images images.to(device) masks masks.to(device).unsqueeze(1) outputs model(images) preds (torch.sigmoid(outputs) 0.5).float() intersection (preds * masks).sum() dice (2.0 * intersection) / (preds.sum() masks.sum() 1e-6) val_dice dice.item() * images.size(0) val_dice / len(val_loader.dataset) if val_dice best_dice: best_dice val_dice torch.save(model.state_dict(), best_resunet.pth) print(fEpoch {epoch1}/{num_epochs}, Loss: {train_loss/len(train_loader.dataset):.4f}, Val Dice: {val_dice:.4f})这段代码里几个点值得注意。masks原始形状如果是(B, H, W)一定要unsqueeze(1)变成单通道的(B,1,H,W)否则和模型输出维度对不上。计算Dice时intersection和preds.sum()是标量除以val_loader.dataset长度得到平均Dice这里我偷懒直接用总像素统计严格做法是先算每个样本的Dice再取平均但对日常验证来说差别不大。关于epoch数量的选择我建议早停加模型保存。医学图像数据量小训练100到200轮完全够用超过300轮后基本就是过拟合了。验证集Dice连续20轮不提升就提前结束训练这也是常规操作。4.3 推理与分割结果可视化训练完成后加载保存的权重在测试图像上推理。推理时要记住一个原则验证和训练的数据预处理必须完全一致否则输入分布变了效果会崩盘。import matplotlib.pyplot as plt def predict_image(model, image_path, device, size256): model.eval() image cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) original_size image.shape image_resized cv2.resize(image, (size, size), interpolationcv2.INTER_LINEAR) image_norm image_resized.astype(np.float32) / 255.0 image_tensor torch.from_numpy(image_norm).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): output model(image_tensor) pred torch.sigmoid(output).cpu().numpy().squeeze() pred (pred 0.5).astype(np.uint8) pred cv2.resize(pred, (original_size[1], original_size[0]), interpolationcv2.INTER_NEAREST) return pred model ResUNet(in_channels1, num_classes1) model.load_state_dict(torch.load(best_resunet.pth, map_locationdevice)) model.to(device) pred predict_image(model, test_image.png, device) plt.figure(figsize(12, 4)) plt.subplot(1, 3, 1) plt.imshow(cv2.imread(test_image.png, cv2.IMREAD_GRAYSCALE), cmapgray) plt.title(原图) plt.axis(off) plt.subplot(1, 3, 2) plt.imshow(pred, cmapgray) plt.title(分割结果) plt.axis(off) plt.subplot(1, 3, 3) plt.imshow(cv2.imread(test_mask.png, cv2.IMREAD_GRAYSCALE), cmapgray) plt.title(真实标注) plt.axis(off) plt.show()推理时最常见的问题是输入图片是三通道彩色图而模型只接受单通道灰度图。要么读取时用cv2.IMREAD_GRAYSCALE要么在预处理里做灰度转换否则会出现维度错误。5. 我踩过的坑训练不稳定与过拟合的排查记录5.1 梯度消失不是唯一问题很多人一听残差网络就以为加了shortcut万事大吉其实训练不稳定的原因可能比想象中复杂。我在训练ResUnet时遇到过loss突然变成NaN的情况。排查过程是第一步检查学习率1e-4通常很安全排除第二步检查数据发现CT图像的HU值范围在[-1000, 3000]之间直接归一化时存在大量负值输入分布太宽导致BN统计量不稳定。我改用窗宽窗位截断后再归一化NaN就消失了。这个坑提醒我医学图像的预处理必须跟任务紧密结合不是随便一个Min-Max归一化就能搞定。另一个常见问题是验证集Dice很高但测试集效果差典型过拟合。医学图像分割数据量小模型很容易把训练集的纹理特征记住而不是语义特征。解决办法有三条一是数据增强尤其弹性形变和随机亮度对比度二是DropBlock或Dropout在解码器部分添加随机丢弃三是减小模型容量把编码器初始通道从32降到16或者减少stage数不要一味追求大模型。5.2 常见问题速查表根据我自己的经历和身边朋友反馈整理了以下高频问题及排查思路大家遇到类似情况可以直接对照。问题现象可能原因解决办法训练loss为NaN学习率太大、输入含有异常值、梯度爆炸调小学习率、检查数据归一化、增加梯度裁剪GPU显存不足batch size过大、输入图像过大减小batch size、使用小图resize、开启梯度累积验证Dice不提升学习率过低、模型欠拟合、标签有错增大学习率、增加训练轮数、检查掩膜分割结果全是背景类别极度不平衡、Dice Loss损失非凸使用混合损失BCEDice、增加前景权重训练快但推理慢模型太大、未开torch.no_grad使用torch.no_grad、尝试半精度推理图像和mask错位数据增强只处理了图像用Albumentations同时处理两者尺寸不匹配报错图像分辨率不是16倍数统一resize到64倍数尺寸这张表不是万能的但能覆盖大部分新手踩坑场景。具体到我做过的一个肝脏分割项目最戏剧化的问题是数据文件夹里混入了几张没有病灶的正常切片标注是全黑mask结果训练后期模型开始把所有图像都预测成背景。把这类“无目标”样本单独处理之后Dice直接从0.72涨到0.85。这种问题常规代码检查很难发现只能靠定期可视化预测结果。6. 经验总结与个性化扩展6.1 这套架构还能用到哪些场景ResUnet不局限于某一个器官凡是二分类或多分类的图像分割任务都可以套用。细胞膜分割、肺结节检测、视网膜血管分割、工业缺陷检测本质上都是像素级分类问题。唯一需要调整的是输入通道数和输出类别数比如RGB医学图像就要把in_channels改成3多器官分割把num_classes改成类别数并用CrossEntropyLoss。我最近在一个视网膜血管分割数据集上做了对比ResUnet的Dice比标准U-Net高约1.5个百分点而且训练曲线更平滑。这个提升主要来自残差结构对边缘细节的保留。血管这类细长结构对低层特征非常敏感残差连接能让浅层细节更顺畅地流向深层解码器。另外如果你觉得普通U-Net对多尺度目标不友好可以考虑在ResUnet的bridge位置加入ASPP模块提取多尺度上下文信息。这个改动对病变区域大小差异显著的场景很有效代码改动量也不大。6.2 最后分享一个实用小技巧训练ResUnet时我习惯在最后几个epoch使用混合精度训练。PyTorch自带的torch.cuda.amp用起来很简单只需要在forward和backward外面包一层autocast和GradScaler显存占用能降低近一半训练速度也有提升。不过要提醒的是混合精度对Dice Loss里的smooth浮点数精度比较敏感如果出现NaN可以把smooth从1e-6改成1e-4或者把损失函数的计算放到float32下。还有一个好用的技巧是训练时保存“最佳模型”和“最后一轮模型”两个权重。最佳模型按验证集Dice保存最后一轮模型也保存下来。有时候测试集上最后一轮模型反而比最佳模型效果好尤其是在训练后期验证集指标震荡的情况下。两个模型在测试集上分别推理取平均预测或者选效果好的那个经常能多拿零点几个点的Dice。说到最后我的体会是网络结构这件事抄成熟方案永远比自己发明靠谱但一定要动手改、动手跑。U-Net是经典ResUnet是进化医学图像分割里真正难的不是网络搭不出来而是数据处理、训练细节和评估逻辑这些“看不见”的环节。把我上面的代码跑通一遍再替换成自己的数据集你会发现自己对分割任务的理解比对着论文看十天还深刻。
做个医学图像分割骨架网络天天在变从早年的FCN到U-Net再到各种Transformer变体但很多新手一上来就抱着U-Net论文死磕模型改了几天换到真实数据上效果还是不行。U-Net当然经典可它的编码器说到底还是纯卷积堆叠网络一旦加深梯度回传和特征复用都会出问题。这绝不是危言耸听我在自己的项目里试过把U-Net编码器从3层加到5层训练loss直接降不下去。后来把ResNet思路融进去改用ResUnet问题一下子清晰了——残差连接不仅让梯度传得更稳特征表达也变得更充分。这篇文章就用PyTorch从头搭一个ResUnet专门用来做医学图像分割环境配置、网络结构、数据处理、训练推理、踩坑记录全部写清楚代码直接可以抄。1. 为什么我放弃“死磕”U-Net转向ResUnet1.1 U-Net的痛点在哪里U-Net最核心的设计是编码器-解码器结构加上跳跃连接。编码器通过逐层下采样提取高层语义信息解码器逐步恢复空间分辨率跳跃连接把编码器各层特征直接拼到解码器对应层弥补下采样丢失的细节。这个思路在二维医学图像分割上确实很有效尤其器官、病灶这类边界相对清晰的目标。可问题在于原始U-Net的编码器基本就是“两次卷积 ReLU 池化”的重复堆叠。网络深度有限一旦尝试加深去提取更抽象的特征梯度在多层卷积之间反复相乘很容易出现梯度消失。更麻烦的是深层卷积层如果学不到有效更新前面浅层学到的特征也会被后续层破坏这就是典型的退化问题。我在处理肝脏CT序列时用标准U-Net训练到30轮之后loss几乎不动加了残差结构后同一套数据和超参数loss还能继续稳定下降这个对比非常直观。另外医学图像本身有几个特点标注数据稀缺、类间不平衡严重、目标边界模糊。比如肿瘤分割中肿瘤区域往往只占整张图像的百分之几如果网络表达能力不够很容易把所有像素都判成背景。ResUnet通过残差连接增强特征复用能力让网络在有限数据下更容易学到判别性特征这也是我从U-Net转向ResUnet的根本原因。1.2 ResUnet的核心思路残差传给谁ResUnet并不是一个非常新的网络但它把两个成熟思路融合得很好一个是ResNet提出的残差学习另一个是U-Net的编码器-解码器结构。残差连接可以做得很细不只在编码器和解码器的每个Block内部加shortcut跳跃连接本身也可以理解为一种残差路径把浅层细节直接“传”给深层解码器。用公式来说普通卷积Block学习的是从输入x到输出y的直接映射而残差Block学习的是残差映射F(x)最终输出为y F(x) x。这样做的价值在于如果某个卷积层对特征的改变没有帮助网络可以直接把权重推向零让输出约等于输入而不会破坏已有特征。这一性质对医学图像分割特别重要因为病灶区域形状差异大、纹理细节丰富网络需要保留更多原始信息。我在实现ResUnet时参考了常见的设计编码器每个stage使用两个卷积层组成一个残差块卷积核大小保持3x3stride1padding1保证特征图尺寸不变下采样使用stride2的卷积或最大池化解码器每个stage先上采样再与对应编码器特征在通道维度拼接然后经过残差块。整体结构看起来就是“U-Net的骨架 ResNet的肌肉”。2. PyTorch环境搭建与工程准备2.1 用Anaconda创建独立环境动手写代码前先把PyTorch环境收拾干净。我强烈建议用Anaconda管理Python环境不要一股脑把包装到系统Python里。医学图像分割项目依赖版本往往互相牵扯今天装这个库、明天升级那个库没有独立环境迟早会出事。打开终端或者Anaconda Prompt执行以下命令创建Python 3.9环境conda create -n resunet python3.9 conda activate resunet为什么推荐Python 3.9因为PyTorch、MONAI、OpenCV、NumPy这些库对3.8到3.10的兼容性都很好3.11以后个别依赖可能会出现编译问题新手没必要在这种地方折腾。创建完环境后顺手把基础库装上pip install numpy opencv-python pillow matplotlib tqdm如果是做医学图像处理后面还需要用到SimpleITK或者MONAI但第一步先别急着全装等用到哪个装哪个避免环境里出现一堆用不上的包。2.2 安装PyTorchGPU版PyTorch安装这个环节看起来简单实际上很多人卡在这里。最稳妥的办法是访问PyTorch官网的Get Started页面根据操作系统和CUDA版本选择命令。这里要提醒一句首先要确认自己电脑的显卡支持CUDA用NVIDIA显卡最省心在命令行执行nvidia-smi就能看到驱动版本和支持的CUDA版本。如果是纯CPU环境学习也可以用CPU版但训练医学图像分割模型会非常慢建议至少用一个老款GTX 1660级别的卡。假设CUDA版本是11.8那么安装命令大致是pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118安装完成后务必验证一下这一步别偷懒import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果cuda.is_available()输出True说明GPU环境正常。很多人的问题出在明明装了CUDA版PyTorch却显示False这时候先检查NVIDIA驱动版本再检查conda环境是否激活最后检查是否装成了CPU版本。我在Win10上用Anaconda加PyCharm做这个项目时就碰到过PyCharm解释器指向了别的环境导致GPU不可用排查了半天才发现是解释器选错了。3. ResUnet网络架构逐层拆解3.1 从Residual Block开始ResUnet的基本单元是残差块我习惯先写一个小类保证编码器和解码器都能复用。代码不长但细节不少。import torch import torch.nn as nn class ResidualBlock(nn.Module): def __init__(self, in_channels, out_channels, stride1): super(ResidualBlock, self).__init__() self.conv1 nn.Conv2d(in_channels, out_channels, kernel_size3, stridestride, padding1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(out_channels, out_channels, kernel_size3, stride1, padding1, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) self.shortcut nn.Sequential() if stride ! 1 or in_channels ! out_channels: self.shortcut nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(out_channels) ) def forward(self, x): identity self.shortcut(x) out self.conv1(x) out self.bn1(out) out self.relu(out) out self.conv2(out) out self.bn2(out) out identity out self.relu(out) return out这里面最关键的是shortcut的设计。当输入和输出的通道数不一致或者特征图尺寸因为stride2被压缩时shortcut必须用1x1卷积调整维度和尺寸否则两个张量无法相加。我一开始图省事直接在shortcut里用nn.MaxPool2d处理下采样后来发现1x1卷积加stride的写法更灵活既能改变通道数也能下采样。3.2 编码器、解码器、跳跃连接完整网络结构可以拆成三个部分编码器、解码器、跳跃连接。编码器负责逐层提取语义信息我用四个stage每层通道数分别是32、64、128、256。这里要说明一下很多人一开始就把通道数设置成64、128、256、512照着ImageNet分类那套来实际在医学图像分割上不一定好因为医学数据量小参数量太大反而容易过拟合。我自己在多个数据集上试下来从32开始逐步翻倍效果和计算量的平衡最好。class ResUNet(nn.Module): def __init__(self, in_channels1, num_classes1): super(ResUNet, self).__init__() self.encoder1 ResidualBlock(in_channels, 32) self.encoder2 ResidualBlock(32, 64, stride2) self.encoder3 ResidualBlock(64, 128, stride2) self.encoder4 ResidualBlock(128, 256, stride2) self.bridge ResidualBlock(256, 512, stride2) self.decoder4 self._make_decoder(512, 256) self.decoder3 self._make_decoder(256, 128) self.decoder2 self._make_decoder(128, 64) self.decoder1 self._make_decoder(64, 32) self.final_conv nn.Conv2d(32, num_classes, kernel_size1) def _make_decoder(self, in_channels, out_channels): return nn.Sequential( nn.ConvTranspose2d(in_channels, out_channels, kernel_size2, stride2), ResidualBlock(out_channels * 2, out_channels) )等等这里有个容易混淆的地方。_make_decoder里我先用转置卷积上采样把通道数从in_channels变成out_channels但之后要和编码器跳跃连接拼接拼接后通道数变成out_channels * 2所以ResidualBlock的输入要设置成out_channels * 2。实际forward里必须先把编码器特征和上采样特征在通道维度拼接再送入残差块。def forward(self, x): e1 self.encoder1(x) # 32 e2 self.encoder2(e1) # 64 e3 self.encoder3(e2) # 128 e4 self.encoder4(e3) # 256 bridge self.bridge(e4) # 512 d4 self.decoder4[0](bridge) d4 torch.cat([d4, e4], dim1) d4 self.decoder4[1](d4) d3 self.decoder3[0](d4) d3 torch.cat([d3, e3], dim1) d3 self.decoder3[1](d3) d2 self.decoder2[0](d3) d2 torch.cat([d2, e2], dim1) d2 self.decoder2[1](d2) d1 self.decoder1[0](d2) d1 torch.cat([d1, e1], dim1) d1 self.decoder1[1](d1) out self.final_conv(d1) return out注意forward里的拼接顺序通常把编码器特征放在前面上采样特征放在后面或者反过来都行但一定要和训练/推理时的数据流保持一致。另外self.decoder4[0]这种写法依赖torch.cat时维度对齐如果输入图像尺寸不是16的倍数转置卷积之后的特征图和编码器特征图尺寸会不一致拼接就会报错。所以数据预处理阶段最好是统一resize到64的倍数比如256x256或者512x512。3.3 损失函数与评估指标为什么这样选医学图像分割里最常用的损失函数是Dice Loss它直接优化Dice系数比逐像素交叉熵更能应对类别不平衡。标准实现如下class DiceLoss(nn.Module): def __init__(self, smooth1e-6): super(DiceLoss, self).__init__() self.smooth smooth def forward(self, pred, target): pred torch.sigmoid(pred) pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() dice (2.0 * intersection self.smooth) / (pred.sum() target.sum() self.smooth) return 1 - dice这里的target是二值掩膜取值为0和1pred经过sigmoid后变成0到1之间的概率值。为什么不直接用交叉熵因为医学图像里前景占比经常低于10%交叉熵的梯度会被背景像素主导网络倾向于把所有像素预测成背景。Dice Loss把前景和背景看作一个整体来度量重叠程度对类别不平衡更鲁棒。评估指标我也习惯额外计算IoU和Dice系数方便和已有论文结果对比。训练时每个epoch结束后在验证集上计算这些指标保存Dice最高的一次权重而不是最后一轮权重这个细节后面还会提到。4. 医学图像分割实战从数据到可视化4.1 数据预处理不要盲目resize医学图像格式五花八门有的是二维PNG切片有的是三维NIfTI文件。这里先用最常见的二维切片场景举例。很多新手拿到数据直接cv2.resize到256x256然后丢进网络结果效果很差。原因在于盲目resize会破坏像素间距比例导致器官或病灶的形状变形。经典的医学图像处理工具SimpleITK或MONAI可以读取NIfTI里的spacing信息按照目标尺寸重采样但二维切片场景下如果标注本身已经是PNG那resize影响倒不大重点要保证所有样本使用相同尺寸。我自己常用的预处理流程是读取图像和对应掩膜统一转成灰度图。对图像做窗宽窗位调整比如CT图像通常关注[-200, 200] HU范围把范围外的值截断。归一化到[0, 1]公式是(img - min) / (max - min)。使用数据增强随机旋转、水平翻转、弹性形变等增强泛化性。如果内存充足可以读取所有数据放到NumPy数组里否则用torch.utils.data.Dataset按需读取。这里有一个关键点数据增强时图像和掩膜必须使用完全相同的变换。用torchvision.transforms里的RandomHorizontalFlip时传入的image和mask是分开处理的很容易出现图像翻转了、mask没翻转的情况导致训练标签错位。我建议用Albumentations库它专门支持同时变换image和mask大幅减少这类低级错误。import albumentations as A train_transform A.Compose([ A.RandomRotate90(), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.1), A.RandomBrightnessContrast(p0.2), ]) def preprocess(image, mask, size256): image cv2.resize(image, (size, size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (size, size), interpolationcv2.INTER_NEAREST) image image.astype(np.float32) / 255.0 mask mask.astype(np.float32) / 255.0 return image, mask注意mask的resize插值方式一定要用INTER_NEAREST如果用线性插值会引入中间灰度值破坏标签的类别语义。4.2 训练流程与关键参数训练一个分割模型参数选择比网络结构本身更影响最终效果。我一般在单卡GPU上使用batch size为8到16图像尺寸256x256编码器初始通道32总参数量约2000万显存占用不高。先把训练集和验证集切出来通常按8:2划分划分时按病人或样本ID分而不是按切片直接随机分避免同一个病人的相邻切片同时出现在训练集和验证集里导致数据泄露。这个细节对医学图像特别重要否则验证集的Dice会虚高真实场景泛化性能很差。优化器我推荐AdamW初始学习率1e-4并配合CosineAnnealingLR或ReduceLROnPlateau。早年我习惯用Adam后来发现AdamW的权重衰减处理更规范配合warmup效果更稳。训练代码如下import torch.optim as optim model ResUNet(in_channels1, num_classes1) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion DiceLoss() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100, eta_min1e-6) best_dice 0.0 num_epochs 100 for epoch in range(num_epochs): model.train() train_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device).unsqueeze(1) # (B,1,H,W) outputs model(images) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() * images.size(0) scheduler.step() # 验证 model.eval() val_dice 0.0 with torch.no_grad(): for images, masks in val_loader: images images.to(device) masks masks.to(device).unsqueeze(1) outputs model(images) preds (torch.sigmoid(outputs) 0.5).float() intersection (preds * masks).sum() dice (2.0 * intersection) / (preds.sum() masks.sum() 1e-6) val_dice dice.item() * images.size(0) val_dice / len(val_loader.dataset) if val_dice best_dice: best_dice val_dice torch.save(model.state_dict(), best_resunet.pth) print(fEpoch {epoch1}/{num_epochs}, Loss: {train_loss/len(train_loader.dataset):.4f}, Val Dice: {val_dice:.4f})这段代码里几个点值得注意。masks原始形状如果是(B, H, W)一定要unsqueeze(1)变成单通道的(B,1,H,W)否则和模型输出维度对不上。计算Dice时intersection和preds.sum()是标量除以val_loader.dataset长度得到平均Dice这里我偷懒直接用总像素统计严格做法是先算每个样本的Dice再取平均但对日常验证来说差别不大。关于epoch数量的选择我建议早停加模型保存。医学图像数据量小训练100到200轮完全够用超过300轮后基本就是过拟合了。验证集Dice连续20轮不提升就提前结束训练这也是常规操作。4.3 推理与分割结果可视化训练完成后加载保存的权重在测试图像上推理。推理时要记住一个原则验证和训练的数据预处理必须完全一致否则输入分布变了效果会崩盘。import matplotlib.pyplot as plt def predict_image(model, image_path, device, size256): model.eval() image cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) original_size image.shape image_resized cv2.resize(image, (size, size), interpolationcv2.INTER_LINEAR) image_norm image_resized.astype(np.float32) / 255.0 image_tensor torch.from_numpy(image_norm).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): output model(image_tensor) pred torch.sigmoid(output).cpu().numpy().squeeze() pred (pred 0.5).astype(np.uint8) pred cv2.resize(pred, (original_size[1], original_size[0]), interpolationcv2.INTER_NEAREST) return pred model ResUNet(in_channels1, num_classes1) model.load_state_dict(torch.load(best_resunet.pth, map_locationdevice)) model.to(device) pred predict_image(model, test_image.png, device) plt.figure(figsize(12, 4)) plt.subplot(1, 3, 1) plt.imshow(cv2.imread(test_image.png, cv2.IMREAD_GRAYSCALE), cmapgray) plt.title(原图) plt.axis(off) plt.subplot(1, 3, 2) plt.imshow(pred, cmapgray) plt.title(分割结果) plt.axis(off) plt.subplot(1, 3, 3) plt.imshow(cv2.imread(test_mask.png, cv2.IMREAD_GRAYSCALE), cmapgray) plt.title(真实标注) plt.axis(off) plt.show()推理时最常见的问题是输入图片是三通道彩色图而模型只接受单通道灰度图。要么读取时用cv2.IMREAD_GRAYSCALE要么在预处理里做灰度转换否则会出现维度错误。5. 我踩过的坑训练不稳定与过拟合的排查记录5.1 梯度消失不是唯一问题很多人一听残差网络就以为加了shortcut万事大吉其实训练不稳定的原因可能比想象中复杂。我在训练ResUnet时遇到过loss突然变成NaN的情况。排查过程是第一步检查学习率1e-4通常很安全排除第二步检查数据发现CT图像的HU值范围在[-1000, 3000]之间直接归一化时存在大量负值输入分布太宽导致BN统计量不稳定。我改用窗宽窗位截断后再归一化NaN就消失了。这个坑提醒我医学图像的预处理必须跟任务紧密结合不是随便一个Min-Max归一化就能搞定。另一个常见问题是验证集Dice很高但测试集效果差典型过拟合。医学图像分割数据量小模型很容易把训练集的纹理特征记住而不是语义特征。解决办法有三条一是数据增强尤其弹性形变和随机亮度对比度二是DropBlock或Dropout在解码器部分添加随机丢弃三是减小模型容量把编码器初始通道从32降到16或者减少stage数不要一味追求大模型。5.2 常见问题速查表根据我自己的经历和身边朋友反馈整理了以下高频问题及排查思路大家遇到类似情况可以直接对照。问题现象可能原因解决办法训练loss为NaN学习率太大、输入含有异常值、梯度爆炸调小学习率、检查数据归一化、增加梯度裁剪GPU显存不足batch size过大、输入图像过大减小batch size、使用小图resize、开启梯度累积验证Dice不提升学习率过低、模型欠拟合、标签有错增大学习率、增加训练轮数、检查掩膜分割结果全是背景类别极度不平衡、Dice Loss损失非凸使用混合损失BCEDice、增加前景权重训练快但推理慢模型太大、未开torch.no_grad使用torch.no_grad、尝试半精度推理图像和mask错位数据增强只处理了图像用Albumentations同时处理两者尺寸不匹配报错图像分辨率不是16倍数统一resize到64倍数尺寸这张表不是万能的但能覆盖大部分新手踩坑场景。具体到我做过的一个肝脏分割项目最戏剧化的问题是数据文件夹里混入了几张没有病灶的正常切片标注是全黑mask结果训练后期模型开始把所有图像都预测成背景。把这类“无目标”样本单独处理之后Dice直接从0.72涨到0.85。这种问题常规代码检查很难发现只能靠定期可视化预测结果。6. 经验总结与个性化扩展6.1 这套架构还能用到哪些场景ResUnet不局限于某一个器官凡是二分类或多分类的图像分割任务都可以套用。细胞膜分割、肺结节检测、视网膜血管分割、工业缺陷检测本质上都是像素级分类问题。唯一需要调整的是输入通道数和输出类别数比如RGB医学图像就要把in_channels改成3多器官分割把num_classes改成类别数并用CrossEntropyLoss。我最近在一个视网膜血管分割数据集上做了对比ResUnet的Dice比标准U-Net高约1.5个百分点而且训练曲线更平滑。这个提升主要来自残差结构对边缘细节的保留。血管这类细长结构对低层特征非常敏感残差连接能让浅层细节更顺畅地流向深层解码器。另外如果你觉得普通U-Net对多尺度目标不友好可以考虑在ResUnet的bridge位置加入ASPP模块提取多尺度上下文信息。这个改动对病变区域大小差异显著的场景很有效代码改动量也不大。6.2 最后分享一个实用小技巧训练ResUnet时我习惯在最后几个epoch使用混合精度训练。PyTorch自带的torch.cuda.amp用起来很简单只需要在forward和backward外面包一层autocast和GradScaler显存占用能降低近一半训练速度也有提升。不过要提醒的是混合精度对Dice Loss里的smooth浮点数精度比较敏感如果出现NaN可以把smooth从1e-6改成1e-4或者把损失函数的计算放到float32下。还有一个好用的技巧是训练时保存“最佳模型”和“最后一轮模型”两个权重。最佳模型按验证集Dice保存最后一轮模型也保存下来。有时候测试集上最后一轮模型反而比最佳模型效果好尤其是在训练后期验证集指标震荡的情况下。两个模型在测试集上分别推理取平均预测或者选效果好的那个经常能多拿零点几个点的Dice。说到最后我的体会是网络结构这件事抄成熟方案永远比自己发明靠谱但一定要动手改、动手跑。U-Net是经典ResUnet是进化医学图像分割里真正难的不是网络搭不出来而是数据处理、训练细节和评估逻辑这些“看不见”的环节。把我上面的代码跑通一遍再替换成自己的数据集你会发现自己对分割任务的理解比对着论文看十天还深刻。
