简介这份资源是面向深度学习入门者与医学图像处理学习者的U-Net脑肿瘤分割完整代码包基于Ronneberger等人提出的经典U-Net架构实现可用于MRI或CT脑部影像中肿瘤区域的自动识别与像素级分割。压缩包共约2000个文件以1989个tif图像数据为主体配合7个Python脚本及少量编译缓存与预览图整体约291MB涵盖数据读取、模型定义、训练与测试等核心模块便于直接复现实验流程。目前已有1025人学习下载。读者可借此理解U-Net下采样与上采样路径的特征融合机制掌握Dice、IoU等分割指标的评估方式并参考20个epoch的训练配置观察模型收敛与过拟合现象适合作为课程设计、毕业设计或医学图像分割入门实践的参考代码。1. 从一张 MRI 到可复现的脑肿瘤分割为什么 U-Net 仍是首选基线脑肿瘤分割这个任务说白了就是把 MRI 影像里的肿瘤区域逐像素标出来——背景、坏死核心、水肿区、增强区每个像素都要归到某一类。临床上这件事靠医生手工勾画一个病例动辄几十分钟不同医生之间的一致性还不高。所以只要做过医学图像方向的人几乎都绕不开 U-Net 脑肿瘤分割完整代码这条路径编码器下采样提特征解码器上采样还原分辨率中间用跳跃连接把浅层的空间细节直接送到深层最后输出一张和原图同尺寸的分割掩膜。它之所以成为基线不是因为最新而是因为稳。数据量小的时候U-Net 的归纳偏置刚好合适训练资源有限的时候它能在一张消费级显卡上跑起来要做对比实验的时候它是最容易被复现的那一档。这篇笔记面向的是想真正把一套完整代码跑通的人——从数据准备、模型搭建、训练循环到评估指标和推理可视化每一步我都会给出可抄的代码和参数说明也会把那些让我翻过车的坑摊开讲。读完你应该能自己搭出一套能出结果的流程而不是停在“看懂原理”这一步。2. 数据准备与预处理把 MRI 变成网络能吃的张量2.1 脑肿瘤分割的数据长什么样常见的脑肿瘤分割数据集一个病例通常包含四种 MRI 模态T1、T1ce增强 T1、T2、FLAIR。四种模态配准到同一空间尺寸一般是 240×240×155 左右。标签是一张整型掩膜取值 0/1/2/3分别对应背景、坏死核心、水肿、增强肿瘤。注意这里有个容易踩的点标签里的 3 个前景类并不是互斥的“三个独立肿瘤”而是层级关系——增强区属于肿瘤核心肿瘤核心又属于整体肿瘤。评估的时候通常按“整体肿瘤”“肿瘤核心”“增强区”三个层级分别算 Dice而不是直接对 0/1/2/3 做四分类 Dice。我一般会把原始数据整理成这样的目录结构后面写 Dataset 的时候直接按这个约定读data/ train/ case_001/ t1.nii.gz t1ce.nii.gz t2.nii.gz flair.nii.gz seg.nii.gz case_002/ ... val/ case_101/ ...2.2 用 nnU-Net 风格做强度归一化MRI 的灰度值没有 CT 那种固定的 HU 单位不同扫描仪、不同序列的强度范围差异很大。直接把原始值喂给网络训练会非常不稳定。常见做法是逐模态做 z-score 归一化并且只在脑组织区域内统计均值和方差——如果连背景的零值一起算进去均值和方差会被拉偏。import numpy as np import nibabel as nib def normalize_modality(volume, brain_maskNone): 对单个模态做 z-score 归一化。 volume: (H, W, D) 的 float 数组 brain_mask: 可选的脑组织掩膜只在掩膜内统计 if brain_mask is not None: values volume[brain_mask 0] else: # 没有掩膜时用非零体素近似脑区 values volume[volume 0] mean values.mean() std values.std() # 防止除零std 太小时给一个下限 std max(std, 1e-8) normalized (volume - mean) / std # 裁掉极端离群值避免个别体素主导梯度 normalized np.clip(normalized, -5.0, 5.0) return normalized.astype(np.float32) def load_case(case_dir): modalities [t1, t1ce, t2, flair] channels [] for mod in modalities: img nib.load(f{case_dir}/{mod}.nii.gz) # 注意 get_fdata 返回 float64转 float32 省显存 vol img.get_fdata().astype(np.float32) vol normalize_modality(vol) channels.append(vol) # 堆成 (4, H, W, D)通道在前方便后续处理 image np.stack(channels, axis0) seg nib.load(f{case_dir}/seg.nii.gz).get_fdata().astype(np.int64) return image, seg这段代码里有两个参数值得说清楚。np.clip的上下限设成 ±5是因为 z-score 之后超过 5 个标准差的体素基本是噪声或伪影留着只会让梯度爆炸。std的下限1e-8是防御性写法实际数据里很少触发但一旦某个模态全是零值没有这行就会得到 NaN然后整个训练 loss 变 NaN排查起来很费时间。2.3 切块策略为什么不能整图直接送进网络240×240×155 的 3D 体积直接送进 3D U-Net显存根本扛不住。标准做法是切 patch。训练时随机采样 patch比如 128×128×128推理时用滑窗带重叠最后对重叠区域做加权平均。切 patch 的时候有个血泪经验不能纯随机采因为肿瘤区域占整个脑体积的比例很小纯随机很可能连续好几个 batch 都采不到前景网络会退化成“全预测背景”。def random_patch(image, seg, patch_size(128, 128, 128), pos_ratio0.6): 以 pos_ratio 的概率从前景附近采样否则全局随机采样。 image: (4, H, W, D) seg: (H, W, D) _, H, W, D image.shape ph, pw, pd patch_size if np.random.rand() pos_ratio: # 找到所有前景体素坐标 fg np.argwhere(seg 0) if len(fg) 0: center fg[np.random.randint(len(fg))] # 以随机前景点为中心加一点抖动 center center np.random.randint(-20, 20, size3) else: center np.array([H // 2, W // 2, D // 2]) else: center np.array([H // 2, W // 2, D // 2]) # 计算 patch 起点并夹到合法范围 start np.clip(center - np.array(patch_size) // 2, 0, None) start np.minimum(start, np.array([H - ph, W - pw, D - pd])) h, w, d start img_patch image[:, h:hph, w:wpw, d:dpd] seg_patch seg[h:hph, w:wpw, d:dpd] return img_patch, seg_patchpos_ratio0.6是我试出来的经验值。设太高比如 0.9网络见到的全是肿瘤附近区域对正常组织的判别力会下降推理时假阳性变多设太低比如 0.2前景样本不够收敛很慢。0.5 到 0.7 之间通常都能用具体看你的数据里肿瘤占比。3. 搭一个能跑通的 3D U-Net结构选择与关键参数3.1 编码器-解码器结构怎么定2D U-Net 处理逐层切片实现简单但丢掉了层间信息对脑肿瘤这种在 Z 方向也有连续性的目标不够友好。3D U-Net 把卷积核换成 3D能同时利用三个方向的空间上下文代价是显存和计算量上去了。我的建议是如果显存小于 12GB先用 2D U-Net 把流程跑通确认数据管线和评估没问题再换 3D。不要一上来就 3D否则数据有问题的时候你分不清是模型的问题还是数据的问题。下面是一个可以直接用的 3D U-Net 实现深度 4 层每层两个 3×3×3 卷积加 InstanceNorm 和 LeakyReLU。为什么用 InstanceNorm 而不是 BatchNorm因为医学图像 patch 的 batch size 通常很小2 到 4BatchNorm 的统计量不稳定InstanceNorm 对每个样本独立归一化小 batch 下更稳。import torch import torch.nn as nn class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv3d(in_ch, out_ch, kernel_size3, padding1, biasFalse), nn.InstanceNorm3d(out_ch, affineTrue), nn.LeakyReLU(0.01, inplaceTrue), nn.Conv3d(out_ch, out_ch, kernel_size3, padding1, biasFalse), nn.InstanceNorm3d(out_ch, affineTrue), nn.LeakyReLU(0.01, inplaceTrue), ) def forward(self, x): return self.block(x) class UNet3D(nn.Module): def __init__(self, in_ch4, num_classes4, 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.MaxPool3d(2) # 瓶颈层 self.bottleneck ConvBlock(base_ch * 8, base_ch * 16) # 解码器转置卷积上采样 跳跃连接拼接 self.up4 nn.ConvTranspose3d(base_ch * 16, base_ch * 8, 2, stride2) self.dec4 ConvBlock(base_ch * 16, base_ch * 8) self.up3 nn.ConvTranspose3d(base_ch * 8, base_ch * 4, 2, stride2) self.dec3 ConvBlock(base_ch * 8, base_ch * 4) self.up2 nn.ConvTranspose3d(base_ch * 4, base_ch * 2, 2, stride2) self.dec2 ConvBlock(base_ch * 4, base_ch * 2) self.up1 nn.ConvTranspose3d(base_ch * 2, base_ch, 2, stride2) self.dec1 ConvBlock(base_ch * 2, base_ch) # 输出层1x1x1 卷积映射到类别数 self.out nn.Conv3d(base_ch, num_classes, kernel_size1) 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)base_ch32是显存和表达力的折中。设成 16模型太小Dice 上不去设成 64patch 128³ 的情况下 12GB 显存基本跑不动。如果你用 2D 版本同样的结构把Conv3d换成Conv2d、MaxPool3d换成MaxPool2d就行base_ch可以开到 64。3.2 损失函数Dice 和交叉熵怎么配脑肿瘤分割里前景和背景极度不平衡纯交叉熵会让网络倾向于全预测背景。常见做法是交叉熵和 Dice 损失加权求和。Dice 直接优化重叠度和评估指标一致交叉熵提供稳定的梯度。权重我一般设成 0.5 比 0.5或者 Dice 稍高一点。class DiceLoss(nn.Module): def __init__(self, num_classes4, smooth1e-5): super().__init__() self.num_classes num_classes self.smooth smooth def forward(self, logits, target): # logits: (B, C, H, W, D), target: (B, H, W, D) probs torch.softmax(logits, dim1) target_onehot torch.zeros_like(probs) target_onehot.scatter_(1, target.unsqueeze(1), 1) # 只在前景类上算 Dice背景类不参与 dice 0.0 for c in range(1, self.num_classes): p probs[:, c] t target_onehot[:, c] inter (p * t).sum() union p.sum() t.sum() dice (2 * inter self.smooth) / (union self.smooth) return 1 - dice / (self.num_classes - 1) # 组合损失 ce nn.CrossEntropyLoss() dice DiceLoss(num_classes4) def criterion(logits, target): return 0.5 * ce(logits, target) 0.5 * dice(logits, target)smooth1e-5是防止分母为零。注意 Dice 损失里我只对前景类求平均背景类不参与——如果把背景也算进去因为背景占比大且容易分对Dice 会虚高掩盖前景分割差的问题。3.3 训练循环与学习率调度优化器用 AdamW初始学习率 1e-3配合余弦退火。训练轮数看数据量几百个病例的话 200 到 300 轮通常够收敛。每个 epoch 结束后在验证集上算 Dice保存最好的那个 checkpoint。from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet3D(in_ch4, num_classes4, base_ch32).to(device) optimizer AdamW(model.parameters(), lr1e-3, weight_decay1e-5) scheduler CosineAnnealingLR(optimizer, T_max200, eta_min1e-6) best_dice 0.0 for epoch in range(200): model.train() for image, seg in train_loader: image, seg image.to(device), seg.to(device) logits model(image) loss criterion(logits, seg) optimizer.zero_grad() loss.backward() # 梯度裁剪防止 3D 网络梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() # 验证 model.eval() val_dice evaluate(model, val_loader, device) if val_dice best_dice: best_dice val_dice torch.save(model.state_dict(), best_unet3d.pth) print(fepoch {epoch}, val_dice {val_dice:.4f}, best {best_dice:.4f})clip_grad_norm_的max_norm1.0是 3D 网络的标配。3D 卷积参数量大偶尔会出现梯度尖峰不裁剪的话 loss 会突然飙到 NaN前面训的全白费。weight_decay1e-5是很轻的正则数据量小的时候可以适当加大到 1e-4。4. 评估与推理Dice 怎么算才和临床对得上4.1 三个层级的 Dice 计算前面提过脑肿瘤分割的评估不是简单的四分类 Dice。标准做法是把标签映射成三个二值问题整体肿瘤标签 123、肿瘤核心标签 13、增强区标签 3。每个层级单独算 Dice最后取平均。def compute_dice(pred, target, num_classes4): pred, target: (H, W, D) 整型数组 返回三个层级的 Dice # 整体肿瘤标签 1, 2, 3 wt_pred (pred 0).astype(np.float32) wt_true (target 0).astype(np.float32) # 肿瘤核心标签 1, 3 tc_pred ((pred 1) | (pred 3)).astype(np.float32) tc_true ((target 1) | (target 3)).astype(np.float32) # 增强区标签 3 et_pred (pred 3).astype(np.float32) et_true (target 3).astype(np.float32) def dice(p, t): inter (p * t).sum() return (2 * inter 1e-5) / (p.sum() t.sum() 1e-5) return dice(wt_pred, wt_true), dice(tc_pred, tc_true), dice(et_pred, et_true)增强区的 Dice 通常最低因为增强区体积小而且不同病例之间差异大。如果你的整体肿瘤 Dice 到了 0.85 但增强区只有 0.6不用太慌这在公开数据集上是很常见的分布。4.2 滑窗推理与重叠加权推理时不能整图送进去要用滑窗。窗口大小和训练 patch 一致步长设为窗口的一半这样相邻窗口有 50% 重叠。重叠区域的预测概率做平均能减少窗口边界的拼接痕迹。def sliding_window_inference(model, image, patch_size(128, 128, 128), num_classes4): image: (4, H, W, D) numpy 数组 返回 (num_classes, H, W, D) 的概率图 model.eval() _, H, W, D image.shape ph, pw, pd patch_size # 步长为窗口的一半 sh, sw, sd ph // 2, pw // 2, pd // 2 prob_map np.zeros((num_classes, H, W, D), dtypenp.float32) count_map np.zeros((H, W, D), dtypenp.float32) with torch.no_grad(): for h in range(0, H - ph 1, sh): for w in range(0, W - pw 1, sw): for d in range(0, D - pd 1, sd): patch image[:, h:hph, w:wpw, d:dpd] patch_t torch.from_numpy(patch).unsqueeze(0).cuda() logits model(patch_t) probs torch.softmax(logits, dim1).squeeze(0).cpu().numpy() prob_map[:, h:hph, w:wpw, d:dpd] probs count_map[h:hph, w:wpw, d:dpd] 1 # 避免除零 count_map np.maximum(count_map, 1) prob_map / count_map[np.newaxis, ...] return prob_map步长设成窗口一半是精度和速度的平衡。步长再小精度提升有限但推理时间翻倍步长等于窗口大小边界处会有明显拼接伪影Dice 会掉几个点。5. 避坑与排查那些让我重训好几次的问题5.1 训练 loss 正常但 Dice 一直是零现象loss 从 1.5 慢慢降到 0.8看起来在收敛但验证集 Dice 始终是 0.00。原因通常是标签映射错了。有些数据集的标签是 0/1/2/3但顺序不是背景/坏死/水肿/增强而是背景/增强/水肿/坏死。如果你按错误的映射算 Dice增强区可能被当成坏死区Dice 自然上不去。解决方法是先打印几个病例的标签唯一值和体素计数确认每个值对应什么再写映射。5.2 显存溢出但 batch size 已经设成 1现象patch 128³、base_ch32、batch size1还是 OOM。原因多半是数据加载时把整个 3D 体积读进了显存或者模型里某处忘了detach。检查 DataLoader 的num_workers和pin_memory设置确保数据在 CPU 上准备好再搬到 GPU。另外如果用了混合精度torch.cuda.amp能省不少显存但要注意 InstanceNorm 在 fp16 下可能不稳定必要时把归一化层强制转成 fp32。5.3 验证集 Dice 比训练集低很多现象训练集 Dice 0.92验证集只有 0.65。这是过拟合的典型信号但脑肿瘤分割里还有一个常见原因是数据泄漏——同一个病例的不同切片被分到了训练集和验证集。正确做法是按病例划分而不是按切片划分。另外如果不同中心的扫描仪参数差异大验证集来自不同中心时 Dice 也会掉这时候要考虑做强度增强或者域适应。5.4 推理结果里有孤立的小块假阳性现象分割结果整体不错但脑子里散布着一些几个体素大的孤立预测块。原因是模型在背景区域也有少量高概率输出。解决方法是推理后做连通域分析去掉体积小于阈值的连通块。阈值一般设 50 到 100 个体素具体看你的分辨率。from scipy import ndimage def remove_small_components(mask, min_size50): mask: (H, W, D) 二值或整型标签 去掉小于 min_size 的连通块 result np.zeros_like(mask) for label in np.unique(mask): if label 0: continue binary (mask label) labeled, num ndimage.label(binary) for i in range(1, num 1): component (labeled i) if component.sum() min_size: result[component] label return result5.5 学习率设大了 loss 直接 NaN现象训练刚开始几个 steploss 就变成 NaN。3D U-Net 对学习率比 2D 敏感1e-3 有时候就太大了。先把学习率降到 1e-4 试如果还 NaN检查输入数据里有没有 NaN 或 Inf——归一化时除以零是常见来源。另外clip_grad_norm_一定要加它能在梯度爆炸的早期把梯度拉回来。6. 从能跑到好用几个让 Dice 再涨几个点的技巧把流程跑通只是第一步。如果你想让结果从“能看”变成“能写进论文”下面这几个技巧是我实际试过有效的。第一个是深度监督。在解码器的每一层输出都接一个辅助分类头算辅助损失最后加权求和。这样浅层也能收到梯度收敛更快最终 Dice 通常能涨 1 到 2 个点。实现上就是在dec4、dec3、dec2后面各加一个 1×1×1 卷积输出类别数然后上采样到原尺寸算损失。辅助损失的权重从 0.1 到 0.3 都行我一般用 0.2。第二个是测试时增强。推理时把输入做几次翻转比如三个轴各翻转一次分别预测再把概率图翻回来平均。这个技巧不改变模型纯推理阶段就能用Dice 通常能涨 0.5 到 1 个点代价是推理时间翻几倍。如果做竞赛或者发论文值得加上。第三个是后处理里的连通域和形态学操作。前面提过去小块其实还可以做一步闭运算把断裂的区域连起来。但要注意形态学操作对增强区这种小目标可能适得其反我一般只对整体肿瘤做闭运算增强区不动。第四个是模型集成。训三到五个不同初始化的 U-Net推理时概率图平均。这是最稳的涨点方式代价是训练和推理成本成倍增加。如果只选一个技巧我会选集成因为它对结果方差的降低最明显。最后说一个我自己的习惯每次改完数据管线或者损失函数先拿一个小数据集比如 10 个病例跑 20 个 epoch看 loss 曲线和 Dice 趋势对不对确认没问题再上全量数据。这样能省下大量等待时间也能更快定位问题。希望帮到你。本文还有配套的精品资源点击获取