ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

U-Net脑肿瘤分割实战:从MRI预处理到3D模型训练与Dice评估

U-Net脑肿瘤分割实战:从MRI预处理到3D模型训练与Dice评估

简介:这份资源是面向深度学习入门者与医学图像处理学习者的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_mask=None): """ 对单个模态做 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, axis=0) 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_ratio=0.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, size=3) 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:h+ph, w:w+pw, d:d+pd] seg_patch = seg[h:h+ph, w:w+pw, d:d+pd] return img_patch, seg_patch

pos_ratio=0.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 到 4),BatchNorm 的统计量不稳定,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_size=3, padding=1, bias=False), nn.InstanceNorm3d(out_ch, affine=True), nn.LeakyReLU(0.01, inplace=True), nn.Conv3d(out_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.InstanceNorm3d(out_ch, affine=True), nn.LeakyReLU(0.01, inplace=True), ) def forward(self, x): return self.block(x) class UNet3D(nn.Module): def __init__(self, in_ch=4, num_classes=4, base_ch=32): 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, stride=2) self.dec4 = ConvBlock(base_ch * 16, base_ch * 8) self.up3 = nn.ConvTranspose3d(base_ch * 8, base_ch * 4, 2, stride=2) self.dec3 = ConvBlock(base_ch * 8, base_ch * 4) self.up2 = nn.ConvTranspose3d(base_ch * 4, base_ch * 2, 2, stride=2) self.dec2 = ConvBlock(base_ch * 4, base_ch * 2) self.up1 = nn.ConvTranspose3d(base_ch * 2, base_ch, 2, stride=2) self.dec1 = ConvBlock(base_ch * 2, base_ch) # 输出层:1x1x1 卷积映射到类别数 self.out = nn.Conv3d(base_ch, num_classes, kernel_size=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], dim=1)) d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1)) d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) return self.out(d1)

base_ch=32是显存和表达力的折中。设成 16,模型太小,Dice 上不去;设成 64,patch 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_classes=4, smooth=1e-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, dim=1) 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_classes=4) def criterion(logits, target): return 0.5 * ce(logits, target) + 0.5 * dice(logits, target)

smooth=1e-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_ch=4, num_classes=4, base_ch=32).to(device) optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5) scheduler = CosineAnnealingLR(optimizer, T_max=200, eta_min=1e-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_norm=1.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(f"epoch {epoch}, val_dice {val_dice:.4f}, best {best_dice:.4f}")

clip_grad_norm_的max_norm=1.0是 3D 网络的标配。3D 卷积参数量大,偶尔会出现梯度尖峰,不裁剪的话 loss 会突然飙到 NaN,前面训的全白费。weight_decay=1e-5是很轻的正则,数据量小的时候可以适当加大到 1e-4。

4. 评估与推理:Dice 怎么算才和临床对得上

4.1 三个层级的 Dice 计算

前面提过,脑肿瘤分割的评估不是简单的四分类 Dice。标准做法是把标签映射成三个二值问题:整体肿瘤(标签 1+2+3)、肿瘤核心(标签 1+3)、增强区(标签 3)。每个层级单独算 Dice,最后取平均。

def compute_dice(pred, target, num_classes=4): """ 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_classes=4): """ 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), dtype=np.float32) count_map = np.zeros((H, W, D), dtype=np.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:h+ph, w:w+pw, d:d+pd] patch_t = torch.from_numpy(patch).unsqueeze(0).cuda() logits = model(patch_t) probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy() prob_map[:, h:h+ph, w:w+pw, d:d+pd] += probs count_map[h:h+ph, w:w+pw, d:d+pd] += 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_ch=32、batch size=1,还是 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_size=50): """ 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 result

5.5 学习率设大了 loss 直接 NaN

现象:训练刚开始几个 step,loss 就变成 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 趋势对不对,确认没问题再上全量数据。这样能省下大量等待时间,也能更快定位问题。希望帮到你。

本文还有配套的精品资源,点击获取

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进