当前位置：首页 > news >正文

从ViT到UNETR：手把手教你用PyTorch和MONAI复现3D医学图像分割SOTA模型

news 2026/6/1 13:37:55

从ViT到UNETR：手把手教你用PyTorch和MONAI复现3D医学图像分割SOTA模型

在医学影像分析领域，3D图像分割一直是极具挑战性的任务。传统的全卷积神经网络（FCNN）虽然在局部特征提取上表现出色，但在捕捉长距离空间依赖关系方面存在明显局限。2021年提出的UNETR模型创新性地将Transformer引入3D医学图像分割，通过序列到序列的建模方式，在BTCV等权威数据集上实现了当时最先进的性能。本文将带您从零开始，使用PyTorch和MONAI框架完整复现这一突破性工作。

1. 环境准备与数据加载

1.1 基础环境配置

首先需要确保开发环境满足以下要求：

# 创建conda环境（推荐） conda create -n unetr python=3.8 conda activate unetr # 安装核心依赖 pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install monai==0.8.1 nibabel==4.0.0

提示：建议使用NVIDIA GPU并安装对应版本的CUDA工具包，3D模型训练对计算资源要求较高

1.2 医学影像数据预处理

医学影像数据通常以NIfTI格式存储，我们需要将其转换为模型可处理的张量格式。BTCV数据集包含30例腹部CT扫描，每例标注了13个器官：

import monai from monai.data import Dataset, DataLoader transforms = monai.transforms.Compose([ monai.transforms.LoadImaged(keys=["image", "label"]), monai.transforms.EnsureChannelFirstd(keys=["image", "label"]), monai.transforms.ScaleIntensityRanged( keys=["image"], a_min=-175, a_max=250, b_min=0.0, b_max=1.0, clip=True), monai.transforms.RandCropByPosNegLabeld( keys=["image", "label"], label_key="label", spatial_size=(96, 96, 96), pos=1, neg=1, num_samples=4, ), ])

2. UNETR核心架构实现

2.1 Transformer编码器模块

UNETR采用ViT-B/16作为基础架构，关键创新在于将3D体数据视为序列处理：

import torch import torch.nn as nn class PatchEmbedding3D(nn.Module): def __init__(self, img_size=96, patch_size=16, in_chans=1, embed_dim=768): super().__init__() self.grid_size = (img_size // patch_size, ) * 3 self.num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2] self.proj = nn.Conv3d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) def forward(self, x): B, C, H, W, D = x.shape x = self.proj(x).flatten(2).transpose(1, 2) # [B, N, embed_dim] return x

2.2 3D CNN解码器设计

解码器通过跳跃连接融合Transformer不同层级的特征：

class UNETRDecoder(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv1 = monai.networks.blocks.Convolution( dimensions=3, in_channels=in_channels, out_channels=out_channels, kernel_size=3, strides=1, norm="instance", act="leakyrelu" ) self.up = monai.networks.blocks.UpSample( dimensions=3, in_channels=out_channels, out_channels=out_channels, scale_factor=2 ) def forward(self, x, skip=None): x = self.conv1(x) if skip is not None: x = torch.cat([x, skip], dim=1) return self.up(x)

3. 关键实现技巧与优化

3.1 内存优化策略

3D Transformer面临的最大挑战是显存消耗，以下是几种有效优化方法：

梯度检查点：在Transformer层中启用梯度检查点
混合精度训练：使用AMP自动混合精度
分块注意力：将大尺寸特征图分块处理

from torch.cuda.amp import autocast def train_step(model, batch): inputs, labels = batch["image"].cuda(), batch["label"].cuda() with autocast(): outputs = model(inputs) loss = dice_loss(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

3.2 位置编码适配

3D位置编码需要特别处理空间维度关系：

class PositionEmbedding3D(nn.Module): def __init__(self, grid_size, embed_dim): super().__init__() self.pos_embed = nn.Parameter( torch.zeros(1, grid_size**3, embed_dim)) def forward(self, x): return x + self.pos_embed

4. 完整训练流程与评估

4.1 训练循环实现

结合MONAI提供的训练工具构建完整流程：

from monai.losses import DiceLoss from monai.metrics import DiceMetric loss_function = DiceLoss(to_onehot_y=True, softmax=True) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) dice_metric = DiceMetric(include_background=False) for epoch in range(200): model.train() for batch in train_loader: train_step(model, batch) model.eval() with torch.no_grad(): for val_batch in val_loader: val_outputs = model(val_batch["image"].cuda()) dice_metric(y_pred=val_outputs, y=val_batch["label"].cuda())

4.2 结果可视化

使用MONAI的可视化工具展示分割效果：

from monai.visualize import plot_2d_or_3d_image plot_2d_or_3d_image( data=val_outputs.argmax(dim=1, keepdim=True), step=0, writer=SummaryWriter(log_dir="logs"), frame_dim=-1, tag="prediction" )