简介基于U-Net与Attention U-Net的医学图像分割系统代码包面向医学影像算法学习者与研究人员解决CT等图像的多类别语义分割与模型对比需求。包内共14个文件含5个Python源码、7个pyc缓存文件、1个requirements依赖清单及1个README说明整体仅16KB轻量且结构清晰。代码覆盖数据处理、模型定义、训练评估与推理预测全流程dataset.py支持自定义路径、随机翻转与CT窗宽窗位对比度增强model.py实现标准U-Net和带注意力门控的Attention U-Net内含卷积块、上采样与循环卷积块便于扩展多分类任务train.py采用AdamW与余弦学习率衰减通过混淆矩阵计算Dice、IoU及各类别精确率、召回率、F1并自动保存最佳模型与JSON训练日志predict.py可加载模型输出原图掩码叠加结果utils.py提供设备检测、指标计算与训练曲线可视化。已有117人学习适合医学图像分割入门、复现经典网络及开展消融实验。1. 小样本医学影像分割为什么默认选U-Net而不是DeepLabV3我接过一个实际项目腹部CT里的肝脏与肿瘤分割标注数据只有47例。当时团队里有人坚持用DeepLabV3因为它在自然图像上刷榜漂亮。结果折腾了两周测试集的Dice只有0.61小肿瘤几乎全部漏检。后来我把骨干换成U-Net同样的训练数据与预处理Dice直接跳到0.83。这件事不是玄学而是模型结构与医学图像特性之间的匹配问题——医学图像分割的核心诉求是小样本、强边界、多尺度目标而U-Net的对称编码器-解码器结构配合跳跃连接天然适合这类任务。Attention U-Net则在这个基础上用注意力门控机制进一步抑制背景噪声专治边界模糊的器官比如胰腺、前列腺。这篇笔记我会从架构拆解、数据预处理、损失函数设计、训练避坑到推理部署完整讲一遍用U-Net和Attention U-Net做医学图像分割的落地路径。无论你是刚接触医学图像分割的研究生还是要在院内设备上部署推理服务的工程岗下面的内容都能直接照着改。所有代码以PyTorch 2.x和SimpleITK为基底兼容Windows和Linux。2. 网络架构拆解U-Net的骨架与Attention U-Net的注意力机制2.1 编码器-解码器与跳跃连接底层设计决定了分割边界U-Net之所以在医学图像分割领域成为默认选择核心在于它的对称U形结构。左侧编码器逐层下采样逐步扩大感受野并提取语义特征右侧解码器逐层上采样把低分辨率的特征图恢复到原始分辨率。真正让U-Net区别于普通FCN的是每一层编码器和解码器之间的跳跃连接Skip Connection。这些跳跃连接把浅层的边缘、纹理细节直接传递给深层解码器弥补了逐层池化带来的空间信息丢失。我在实际项目中体会最深的一点是医学图像CT、MRI、超声的边界往往没有自然图像那么锐利器官与周围组织的灰度对比度可能只有十几个HU值。如果只用深层的语义特征去做分割边界基本会糊成一团。跳跃连接相当于给解码器每层都配了一份“高分辨率底图”让模型在恢复分辨率时有细节可依。在实现上标准的U-Net一般做4次下采样每次下采样特征图尺寸减半、通道数翻倍。初始通道数常见为32或64我个人的习惯是CT体数据用32起步MRI用64起步因为MRI的纹理细节更丰富需要更宽的浅层特征通道。下面是一个直接可用的U-Net编码器-解码器骨架没有依赖第三方分割库方便你看清每一层在干什么import torch import torch.nn as nn class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch1, out_ch3, base_ch32): super().__init__() # 编码器每次下采样后通道数翻倍特征图尺寸减半 self.enc1 ConvBlock(in_ch, base_ch) self.enc2 ConvBlock(base_ch, base_ch * 2) self.enc3 ConvBlock(base_ch * 2, base_ch * 4) self.enc4 ConvBlock(base_ch * 4, base_ch * 8) self.pool nn.MaxPool2d(2) # 瓶颈层 self.bottleneck ConvBlock(base_ch * 8, base_ch * 16) # 解码器每次上采样后与对应编码层特征拼接 self.up4 nn.ConvTranspose2d(base_ch * 16, base_ch * 8, 2, stride2) self.dec4 ConvBlock(base_ch * 16, base_ch * 8) self.up3 nn.ConvTranspose2d(base_ch * 8, base_ch * 4, 2, stride2) self.dec3 ConvBlock(base_ch * 8, base_ch * 4) self.up2 nn.ConvTranspose2d(base_ch * 4, base_ch * 2, 2, stride2) self.dec2 ConvBlock(base_ch * 4, base_ch * 2) self.up1 nn.ConvTranspose2d(base_ch * 2, base_ch, 2, stride2) self.dec1 ConvBlock(base_ch * 2, base_ch) self.out nn.Conv2d(base_ch, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)这段代码里有三个关键参数要注意。in_ch是输入图像的通道数灰度CT和MRI是1如果你做了多模态融合比如PET-CT叠加就需要改成2或更多。out_ch是分割类别数比如肝脏与肿瘤两类的任务加上背景就是3。base_ch控制整个网络的宽度直接决定参数量与显存占用base_ch从64提到96Dice可能会涨0.02左右但显存占用几乎翻倍。ConvBlock里每个卷积层后面都跟了BatchNorm和ReLU这是当前训练稳定性的基本配置。如果你用了大批量训练batch size大于8BatchNorm没问题但如果是小批量只有2或4建议换成GroupNorm否则验证集上的指标会随batch变化而抖动。2.2 Attention Gate在哪里插入参数、计算量与时序问题Attention U-Net在U-Net基础上引入了注意力门控Attention Gate简称AG核心思想是让解码器在融合编码器特征之前先判断哪些空间位置真正需要被关注。传统U-Net是把所有跳跃连接的特征无差别拼接进来但背景区域比如CT图像里的床板、空气、脂肪会占用大量特征通道。Attention Gate通过一个额外生成的门控信号对跳跃连接的特征图做空间权重调制让模型聚焦在目标器官附近。Attention Gate的插入位置很有讲究。原始论文是在每层解码器的拼接操作之前加一个AG但我在实践里发现最浅层的AG对应第一层分辨率最高收益最小因为浅层特征本身就以边缘信息为主噪声也是逐像素的门控信号难以有效区分。我一般的做法是只对第二到第四层的跳跃连接加AG既能减少显存占用又能拿到和完整版本接近的效果。Attention Gate的计算流程如下编码器特征x跳跃连接的特征图与门控信号g来自下一层解码器上采样后的输出分别经过1x1卷积调整通道数相加后经过ReLU和1x1卷积再用Sigmoid生成[0,1]之间的注意力系数最后与原特征x逐元素相乘。下面是一个在PyTorch里实际跑通的Attention Gate模块class AttentionGate(nn.Module): def __init__(self, in_ch, gate_ch, inter_chNone): super().__init__() inter_ch inter_ch or in_ch // 2 # 对跳跃连接的特征做通道变换 self.conv_x nn.Conv2d(in_ch, inter_ch, 1) # 对门控信号做通道变换gate_ch来自解码器上采样层 self.conv_g nn.Conv2d(gate_ch, inter_ch, 1) # 融合后生成注意力系数 self.psi nn.Sequential( nn.ReLU(inplaceTrue), nn.Conv2d(inter_ch, 1, 1), nn.Sigmoid() ) def forward(self, x, g): x_trans self.conv_x(x) g_trans self.conv_g(g) # 两个特征图尺寸必须一致否则需要插值 if x_trans.shape ! g_trans.shape: g_trans nn.functional.interpolate( g_trans, sizex_trans.shape[2:], modebilinear) attn self.psi(x_trans g_trans) return x * attn使用这个模块时in_ch是编码器特征图的通道数gate_ch是解码器上采样特征图的通道数。有个容易踩坑的地方在标准U-Net里第n层解码器的上采样输出通道数等于第n层编码器输出通道数但经过torch.cat拼接后送入下一层ConvBlock的通道数会翻倍所以Attention Gate的gate_ch应该是上一层ConvBlock的输出通道数而不是拼接后的通道数。Attention U-Net相对标准U-Net的参数量增加大约8%到12%这个代价换来的收益在边界模糊的器官上非常明显。我拿胰腺分割做过对比在完全相同的训练配置下Attention U-Net的Dice比标准U-Net高约0.03而肝脏这种边界清晰的器官两者几乎持平。所以如果你要分割的目标是胰腺、前列腺、肾脏内部肿瘤这类低对比度结构我建议直接上Attention U-Net。3. 医学图像分割数据准备NIfTI切片、窗宽窗位与标签重编码3.1 用SimpleITK把NIfTI和DICOM转成2D切片医学图像分割的落地流程里数据处理占的时间比重往往超过模型训练本身。一份典型的CT数据是NIfTI格式.nii.gz三维数组的shape可能是[512, 512, 300]其中前两个维度是横断面分辨率第三个维度是切片数。大多数分割模型是2D的所以第一步就是把三维体数据切成2D切片。这里有一个关键点NIfTI文件自带方向信息Direction矩阵同一个病人的数据在不同设备上扫描存储的轴向顺序可能不一样。如果不做统一处理有的数据切出来是横断面有的是冠状面模型训练会严重震荡。我一般用SimpleITK读取后先重采样到各向同性体素比如1mm x 1mm x 1mm再统一以轴向Axial方向切片。import SimpleITK as sitk import numpy as np def load_and_resample_nifti(image_path, target_spacing(1.0, 1.0, 1.0)): 读取NIfTI重采样到目标体素间距返回numpy数组 img sitk.ReadImage(image_path) # 获取原始体素间距 original_spacing img.GetSpacing() original_size img.GetSize() # 计算新尺寸原尺寸 * 原间距 / 目标间距 new_size [ int(round(original_size[0] * original_spacing[0] / target_spacing[0])), int(round(original_size[1] * original_spacing[1] / target_spacing[1])), int(round(original_size[2] * original_spacing[2] / target_spacing[2])) ] # Resample滤波器默认用线性插值 resampler sitk.ResampleImageFilter() resampler.SetSize(new_size) resampler.SetOutputSpacing(target_spacing) resampler.SetOutputDirection(img.GetDirection()) resampler.SetOutputOrigin(img.GetOrigin()) resampler.SetInterpolator(sitk.sitkLinear) img_resampled resampler.Execute(img) return sitk.GetArrayFromImage(img_resampled), img_resampled.GetDirection() def extract_slices(volume, label, target_slice_axis0): 沿指定轴切片。SimpleITK的GetArrayFromImage返回的数组是(z, y, x)顺序 # 默认target_slice_axis0表示沿z轴轴向切片 num_slices volume.shape[target_slice_axis] for i in range(num_slices): if target_slice_axis 0: img_slice volume[i, :, :] label_slice label[i, :, :] # 过滤掉没有标签的切片 if np.max(label_slice) 0: yield img_slice, label_slice这段代码里target_spacing(1.0, 1.0, 1.0)的意思是重采样到1毫米等方性体素。为什么要做这个步骤因为不同CT扫描仪的层厚可能是1.25mm、2.5mm甚至5mm如果不统一同一个器官在不同样本里的形态会被拉伸或压缩模型学到的形状特征会互相矛盾。切片方向的选择上腹部和胸部CT我会优先用轴向即z轴切片因为CT扫描本身就是轴向采集的层内分辨率最高层间分辨率相对低。脑部MRI则三种方向都有如果训练数据量不够可以三种方向都切天然做了数据增强。extract_slices函数里的判断np.max(label_slice) 0是过滤空白切片这一步能把训练数据量减少20%到40%同时避免模型被大量背景切片带偏。3.2 归一化、窗宽窗位与数据增强的参数怎么设医学图像的像素值范围与自然图像完全不同。CT值单位是HUHounsfield Unit理论上范围是-1024到3071但人体软组织实际集中在-200到200之间。如果直接做min-max归一化大量组织在数值上会被压缩到接近0对比度完全丢失。正确做法是先做窗宽窗位截断再归一化。不同器官的窗宽窗位差异很大肝脏一般窗宽400、窗位40肺窗窗宽1500、窗位-600脑组织窗宽80、窗位40。我处理腹部CT时的默认配置是def ct_window_normalize(image, window_width400, window_level40): CT窗宽窗位截断 min-max归一化到[0,1] lower window_level - window_width / 2 upper window_level window_width / 2 # 先截断到窗宽范围 image_clipped np.clip(image, lower, upper) # 再线性拉伸到[0, 1] image_norm (image_clipped - lower) / (upper - lower) return image_norm.astype(np.float32)window_width400, window_level40是我做肝脏分割的默认参数如果你用的公开数据集比如LiTS或CHAOS已经预处理过了这段代码可以跳过。MRI数据没有标准的HU值范围不同序列T1、T2的信号强度含义不同我一般不做窗宽截断直接用z-score归一化(image - mean) / std统计范围取每例数据自身或整个训练集的均值方差。数据增强方面医学图像和自然图像有三个重要的区别。第一不能做水平翻转和垂直翻转之外的随机旋转大角度因为器官有固定的解剖朝向第二弹性形变是最有效的增强方式能模拟器官在不同病人体内的形态差异第三亮度对比度扰动不能太大CT值本身有物理含义过度扰动会破坏组织对比度。我用的是albumentations库配置如下import albumentations as A train_transform A.Compose([ A.RandomRotate90(p0.3), A.Flip(p0.5), A.ElasticTransform(alpha3, sigma50, p0.3), A.RandomBrightnessContrast(brightness_limit0.1, contrast_limit0.1, p0.3), ]) # 测试集只做最基本的预处理不做增强 test_transform A.Compose([ A.NoOp() ])ElasticTransform的alpha3, sigma50是我调过多个数据集后比较稳的参数组合。alpha太大形变会过于剧烈导致器官扭曲成不合理的形状太小则没有增强效果。如果你要分割的是本身形态就很多变的器官比如胃、小肠alpha可以调到5-6但要注意标签的形变与图像保持一致albumentations在这方面处理得很好。4. 用PyTorch从零训练损失函数、评估指标与超参数设置4.1 Dice Loss与多类交叉熵的组合策略医学图像分割里最经典的损失函数是Dice Loss它直接优化分割结果与真实标签的重叠度公式上就是1减去Dice系数。但Dice Loss有个实际问题在训练初期模型输出概率分布还接近均匀梯度信号不稳定Loss会剧烈跳动。我在项目里通常把Dice Loss与交叉熵按比例混合用交叉熵的稳定梯度做“引导”用Dice Loss做“精修”。如果是多类别分割比如肝脏、肝肿瘤、肝静脉三类直接计算多类Dice会导致小体积类别被大体积类别淹没。我的做法是分别计算每个类别的Dice再按类别体积比例的倒数做加权。下面是我在实际训练中使用的组合损失函数class CombinedLoss(nn.Module): def __init__(self, num_classes, class_weightsNone, dice_weight0.7): super().__init__() self.num_classes num_classes self.dice_weight dice_weight # class_weights: 每个类别在Dice计算中的权重默认全1 self.class_weights class_weights or [1.0] * num_classes self.ce nn.CrossEntropyLoss() def dice_loss(self, pred, target, eps1e-6): # pred: [B, C, H, W] 经过softmax的概率 # target: [B, H, W] 整数标签 dice_total 0.0 pred torch.softmax(pred, dim1) for c in range(self.num_classes): p pred[:, c] t (target c).float() intersection (p * t).sum() union p.sum() t.sum() dice (2 * intersection eps) / (union eps) dice_total (1 - dice) * self.class_weights[c] return dice_total / self.num_classes def forward(self, pred, target): dice self.dice_loss(pred, target) ce self.ce(pred, target) return self.dice_weight * dice (1 - self.dice_weight) * cedice_weight0.7表示Dice Loss占总损失的70%这个比例在多个数据集上表现比较稳定。class_weights的赋值有个实用技巧先统计训练集每个类别的像素占比然后取倒数并归一化。比如肝脏占30%、肿瘤占2%、背景占68%权重就设为[1.0/0.68, 1.0/0.30, 1.0/0.02]再归一化。这个做法能显著提升小目标的召回率。还有一个需要注意的细节在计算Dice时预测概率与标签都用浮点数计算不要用one-hot编码后再做矩阵乘法那样显存占用会高出一大截。上面的写法直接通过pred[:, c]取概率通道配合布尔比较生成目标掩膜显存效率更高。4.2 学习率、batch size与训练轮次一张可抄的参数表医学图像分割的训练配置与自然图像分类有很大不同。首先batch size受限于显存一般不可能用到64或128其次分割任务需要更多轮次才能稳定收敛因为Dice Loss的梯度信噪比对学习率非常敏感。我用一张表给出我实际验证过的起始配置基于单张NVIDIA RTX 3090或A500024GB显存参数项推荐值说明输入分辨率256x256 或 288x288太小丢失边界细节太大显存吃紧batch size8256x256或 4288x288再小建议用GroupNorm替代BatchNorm初始学习率1e-4AdamW不要学自然图像用1e-3Dice会飞学习率调度余弦退火最小1e-6比ReduceLROnPlateau更稳训练轮次100-200监控验证Dice连续20轮不涨就停优化器AdamWweight_decay1e-4Adam不带W也行但W能压过拟合梯度裁剪最大范数1.0防止Dice Loss的异常梯度这里重点说学习率。医学图像分割的公理是学习率大了必炸。我见过太多人在训练初期看到Loss降得很快结果第30轮时验证Dice突然从0.8跌到0.2这就是学习率过大导致模型跳出了之前的优化区域。我现在的做法是固定使用余弦退火初始学习率1e-4最低学习率1e-6前10轮用线性warmup从1e-5升到1e-4。下面是一个基于上面配置的完整训练循环骨架包含验证指标计算和模型保存import torch from torch.cuda.amp import autocast, GradScaler def train_one_epoch(model, loader, optimizer, scaler, criterion): model.train() epoch_loss 0.0 for images, labels in loader: images images.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) optimizer.zero_grad() with autocast(): logits model(images) loss criterion(logits, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() epoch_loss loss.item() return epoch_loss / len(loader) def validate(model, loader, criterion): 验证集只计算Dice系数返回每类Dice均值 model.eval() dice_scores [] for images, labels in loader: images images.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) with torch.no_grad(): logits model(images) pred torch.argmax(torch.softmax(logits, dim1), dim1) # 逐类别计算Dice num_classes logits.shape[1] for c in range(num_classes): p (pred c).float() t (labels c).float() inter (p * t).sum() union p.sum() t.sum() dice 2 * inter / union.clamp(min1e-6) dice_scores.append(dice.item()) return sum(dice_scores) / len(dice_scores)GradScaler配合autocast是PyTorch的混合精度训练能把训练速度提升40%到60%。如果你用的是A100或更新的卡还可以把torch.cuda.amp换成torch.autocast(device_typecuda)API兼容性更好。训练过程中我建议在验证集上同时记录Dice和IoU两个指标如果Dice在涨但IoU在降通常是模型输出的边界过宽或过窄需要检查损失权重和标签质量。5. 训练避坑记录小目标丢失、类别不平衡与显存溢出5.1 现象小目标器官在预测结果里消失了这是我做肝肿瘤分割时第一个遇到的坑。训练集Dice正常升高验证集肝脏Dice到了0.9但肿瘤Dice始终在0.1到0.2之间徘徊预测结果里肿瘤区域几乎全是背景类。原因有两层。第一肿瘤在整张CT切片中占比极低平均不到2%交叉熵损失被背景像素主导模型倾向于把所有像素都判为背景第二Dice Loss在训练初期对小目标的梯度不稳定几次异常梯度就把小目标特征通道压制了。解决方法是组合拳。首先给类别的Dice加权肿瘤权重设为肝脏的3到5倍让模型在优化时更重视肿瘤区域其次把输入分辨率从256提高到288或320小目标在更高分辨率下能提供更多有效像素参与训练最后把训练轮次拉长小目标通常要在第60轮之后才稳定抬头不要提前停止。如果你用的是Attention U-Net再加上针对目标区域的注意力监督收敛会更快。5.2 现象Dice Loss在冷启动阶段Loss值异常跳动现象是训练前几步Loss值在2.0到5.0之间剧烈震荡甚至出现NaN。原因很明确Dice Loss的梯度与小目标的分母强相关当某个类别的预测概率接近于0时分母趋近于真实目标的像素总数梯度会异常巨大超过数值稳定范围。我的解决方法是三件事同时做。第一初始化后先单独跑5轮的交叉熵损失把模型权重拉到一个合理区域再切换到联合损失第二把学习率从默认的1e-4降到3e-5作为冷启动第三在Dice计算中加入更大的平滑项eps从1e-6提高到1e-4。另外模型最后一层要确保输出不做Sigmoid而是走Logits直出Dice Loss内部再做Softmax这个顺序不要写反。5.3 现象切片方向不一致导致验证指标虚高有一次我换了一批新数据没有做方向统一就直接训练结果验证Dice高达0.87但模型在真实测试集上只有0.55。排查后发现新数据的方向矩阵与训练集不一致导致同一器官被切成了不同视角的切片模型在验证集上看到的形态和训练分布相近指标自然虚高。解决办法在数据预处理阶段就要固定所有NIfTI文件加载后用SimpleITK的Resample统一到同一个Direction矩阵同时把所有图像和标签按同一方向重采样再进行切片。如果你拿到的是DICOM序列建议先用dcm2niix转成NIfTI再走同一套流程。这个坑不报错、不闪退隐蔽性很强但后果是整个模型不可用。5.4 现象显存爆掉batch size根本拉不起来训练刚开始就报CUDA out of memory最常见的原因不是batch size太大而是输入分辨率太高或者模型通道数设得太宽。我曾经在一个3D医学图像数据集上尝试直接使用完整分辨率切片结果24GB显存连batch size为2都跑不动。这里给出三个可选方案按我的优先级排序。第一降低输入分辨率到224或256通常Dice损失在0.01以内第二把base_ch从64降到48或32参数量减少一半多效果损失可控第三使用梯度累积模拟更大的batch size在optimizer.step()之前做4到8次前向反向累积代码层面只要修改训练循环把梯度累积变量累加后再清零。如果你确实需要高分辨率输入我建议直接用patch-based训练从512x512的切片里随机裁剪256x256的patch既能保证足够的上下文又能把batch size维持在8。这是目前医学图像分割训练的主流做法。6. 把模型从验证集推到现场滑窗推理与多模型集成训练指标好看不等于能上线。医学图像是三维体数据推理时不可能把整卷CT直接塞进2D模型滑窗推理是必须实现的工程环节。我的做法是先按训练时的切片方向把三维数据切好逐张预测再用重叠滑窗和权重融合减少切片边缘的不连续性。滑窗参数上窗口大小保持与训练输入一致比如256x256重叠率设为50%。每一张切片预测时覆盖两种情况如果目标器官尺寸变化不大直接小patch预测如果遇到大器官如肝脏用窗口滑动拼接重叠区域取两个窗口预测概率的平均值而不是硬投票概率平均能保留模型的置信度信息。def sliding_window_inference(volume, model, window_size256, overlap0.5): 对3D体数据做滑窗推理返回与输入同尺寸的概率图 _, h, w volume.shape stride int(window_size * (1 - overlap)) # 初始化概率累积和计数矩阵 prob_map np.zeros((num_classes, h, w)) count_map np.zeros((h, w)) for i in range(0, h - window_size 1, stride): for j in range(0, w - window_size 1, stride): patch volume[:, i:i window_size, j:j window_size] patch_input torch.from_numpy(patch).unsqueeze(0).float().cuda() with torch.no_grad(): logits model(patch_input) prob torch.softmax(logits, dim1)[0].cpu().numpy() # 累加预测概率 prob_map[:, i:i window_size, j:j window_size] prob count_map[i:i window_size, j:j window_size] 1 # 归一化 count_map np.clip(count_map, a_min1, a_maxNone) prob_map prob_map / count_map return prob_map推理阶段还有一个容易忽略的点输入图像的归一化必须与训练时完全一致。如果训练时用了窗宽窗位归一化推理时也必须用同一套参数不能直接用原始HU值送进模型。最后说多模型集成。我常用的方案是训练两个模型一个标准U-Net一个Attention U-Net然后在推理时对两者的Softmax概率取平均。实验数据表明这种集成方式比单独使用Attention U-Net提升约1到2个百分点的Dice尤其在边界区域效果显著。代价是推理时间翻倍但如果你的场景对实时性要求不高比如离线辅助诊断这个代价是值得的。把模型真正推到现场之前一定要回看一遍训练时最差的预测样本——那些Dice垫底的验证集切片往往藏着你没预料到的坑。我在这个环节吃过亏模型整体指标达标但某一种少见肿瘤形态几乎全错后来发现是训练集里该类形态的样本太少。这个教训现在变成了我的固定习惯每次训练结束都输出最差的10个样本让医生复核一遍。希望帮到你。本文还有配套的精品资源点击获取