从零构建Swin-UNet：Transformer医学图像分割实战指南

当医学影像遇上Transformer架构，一场关于精准分割的技术革命正在发生。传统CNN在CT、MRI等二维医学图像处理中虽表现稳定，但面对器官边缘模糊、病灶形态多变等复杂场景时，其局部感受野的局限性逐渐显现。本文将带您用PyTorch 1.7完整实现Swin-UNet——这个将Swin Transformer与UNet架构完美融合的标杆模型，从环境搭建到模型部署，破解医学图像分割中的维度对齐、长程依赖等核心难题。

1. 环境配置与数据准备

1.1 开发环境搭建

推荐使用Anaconda创建专属Python 3.6环境，避免依赖冲突：

bash复制conda create -n swin_unet python=3.6
conda activate swin_unet
pip install torch==1.7.0+cu110 torchvision==0.8.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html
pip install opencv-python nibabel scikit-image

关键组件版本对照表：

提示：若使用Colab环境，需在Notebook开头添加!pip install --upgrade torch确保版本兼容

1.2 医学图像预处理

医学影像数据通常以DICOM或NIfTI格式存储，需要特殊处理：

python复制import nibabel as nib
import cv2

def load_nifti(path):
    img = nib.load(path).get_fdata()
    # 归一化到[0,255]并转换为uint8
    img = (img - img.min()) / (img.max() - img.min()) * 255
    return img.astype('uint8')

def preprocess(img, target_size=224):
    # 多模态图像取首通道
    if len(img.shape) > 2:
        img = img[..., 0]
    # 等比例缩放+边缘填充
    h, w = img.shape
    scale = target_size / max(h, w)
    img = cv2.resize(img, (int(w*scale), int(h*scale)))
    pad_h = target_size - img.shape[0]
    pad_w = target_size - img.shape[1]
    img = cv2.copyMakeBorder(img, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)
    return img

典型医学数据集处理流程：

1. ACDC心脏MRI：包含右心室(RV)、左心室(LV)、心肌(Myo)三部分标注
2. Synapse多器官CT：涉及主动脉、胆囊等8个腹部器官
3. BraTS脑肿瘤：多模态MRI的肿瘤区域分割

2. Swin-UNet核心模块实现

2.1 Patch Embedding层

将图像转换为Transformer可处理的token序列：

python复制import torch.nn as nn

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=4, in_chans=1, embed_dim=96):
        super().__init__()
        self.img_size = (img_size, img_size)
        self.patch_size = (patch_size, patch_size)
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_chans, embed_dim, 
                             kernel_size=patch_size, 
                             stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # [B, num_patches, embed_dim]
        return x

维度变换过程图解：

code复制输入: [B, 1, 224, 224]
↓ Conv2d(kernel=4, stride=4)
中间: [B, 96, 56, 56] 
↓ flatten+transpose
输出: [B, 3136, 96]  # 3136=56*56

2.2 Swin Transformer Block

实现带窗口偏移的多头自注意力机制：

python复制class SwinTransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, window_size=7, shift_size=0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = WindowAttention(
            dim, window_size=(window_size, window_size), num_heads=num_heads)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim))
        
        self.window_size = window_size
        self.shift_size = shift_size

    def forward(self, x):
        H, W = self.H, self.W
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"
        
        shortcut = x
        x = self.norm1(x)
        x = x.view(B, H, W, C)
        
        # 窗口划分
        if self.shift_size > 0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_x = x
            
        x_windows = window_partition(shifted_x, self.window_size)
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
        
        # 窗口注意力
        attn_windows = self.attn(x_windows)
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        
        # 窗口合并
        shifted_x = window_reverse(attn_windows, self.window_size, H, W)
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            x = shifted_x
        x = x.view(B, H * W, C)
        
        # FFN
        x = shortcut + x
        x = x + self.mlp(self.norm2(x))
        return x

注意：WindowAttention需实现相对位置编码，完整代码见配套GitHub仓库

2.3 Patch Expanding上采样

替代传统反卷积的Transformer友好方案：

python复制class PatchExpanding(nn.Module):
    def __init__(self, input_resolution, dim):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.expand = nn.Linear(dim, 2*dim, bias=False)
        self.norm = nn.LayerNorm(dim // 2)

    def forward(self, x):
        H, W = self.input_resolution
        x = self.expand(x)
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"
        
        x = x.view(B, H, W, C)
        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', 
                     p1=2, p2=2, c=C//4)
        x = x.view(B, -1, C//4)
        x = self.norm(x)
        return x

上采样过程维度变化：

code复制输入: [B, 56*56, 384]
↓ 线性扩展
中间: [B, 3136, 768] 
↓ 像素重排
输出: [B, 112*112, 192]

3. 完整模型架构与训练技巧

3.1 对称编解码结构搭建

python复制class SwinUNet(nn.Module):
    def __init__(self, img_size=224, in_chans=1, num_classes=3):
        super().__init__()
        depths = [2, 2, 6, 2]
        num_heads = [3, 6, 12, 24]
        embed_dim = 96
        
        # Encoder
        self.patch_embed = PatchEmbed(img_size, 4, in_chans, embed_dim)
        self.encoder_layers = nn.ModuleList([
            BasicLayer(dim=embed_dim*2**i,
                      depth=depths[i],
                      num_heads=num_heads[i],
                      downsample=PatchMerging if i < 3 else None)
            for i in range(4)])
        
        # Bottleneck
        self.bottleneck = BasicLayer(dim=embed_dim*8,
                                   depth=depths[-1],
                                   num_heads=num_heads[-1])
        
        # Decoder
        self.decoder_layers = nn.ModuleList([
            BasicLayer(dim=embed_dim*2**(3-i),
                      depth=depths[3-i],
                      num_heads=num_heads[3-i],
                      upsample=PatchExpanding)
            for i in range(4)])
        
        # Segmentation head
        self.head = nn.Conv2d(embed_dim, num_classes, kernel_size=1)

    def forward(self, x):
        # Encoder
        x = self.patch_embed(x)
        enc_features = []
        for layer in self.encoder_layers:
            x = layer(x)
            enc_features.append(x)
        
        # Bottleneck
        x = self.bottleneck(x)
        
        # Decoder
        for i, layer in enumerate(self.decoder_layers):
            x = layer(torch.cat([x, enc_features[3-i]], dim=-1))
        
        # 恢复空间维度
        B, L, C = x.shape
        H = W = int(L ** 0.5)
        x = x.transpose(1, 2).view(B, C, H, W)
        return self.head(x)

模型参数量统计（输入224×224）：

3.2 预训练权重加载

使用ImageNet预训练参数初始化：

python复制def load_pretrained(model, checkpoint_path):
    state_dict = torch.load(checkpoint_path)['model']
    
    # 过滤decoder相关键值
    pretrained_dict = {k: v for k, v in state_dict.items() 
                      if 'decoder' not in k and 'head' not in k}
    
    # 修改encoder部分键名
    new_dict = {}
    for k, v in pretrained_dict.items():
        if 'layers' in k:
            new_k = k.replace('layers', 'encoder_layers')
            new_dict[new_k] = v
        else:
            new_dict[k] = v
    
    model.load_state_dict(new_dict, strict=False)

提示：官方Swin-T预训练模型可在HuggingFace获取

3.3 损失函数与评估指标

医学分割常用复合损失：

python复制class DiceBCELoss(nn.Module):
    def __init__(self, smooth=1e-5):
        super().__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        pred = torch.sigmoid(pred)
        intersection = (pred * target).sum(dim=(2,3))
        union = pred.sum(dim=(2,3)) + target.sum(dim=(2,3))
        dice = (2.*intersection + self.smooth)/(union + self.smooth)
        bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none').mean(dim=(1,2))
        return (1 - dice).mean() + bce.mean()

评估指标实现：

python复制def dice_score(pred, target):
    pred = (pred > 0.5).float()
    intersection = (pred * target).sum()
    union = pred.sum() + target.sum()
    return 2 * intersection / (union + 1e-8)

def hausdorff_distance(pred, target):
    pred = pred.cpu().numpy()
    target = target.cpu().numpy()
    return max(
        directed_hausdorff(pred, target)[0],
        directed_hausdorff(target, pred)[0]
    )

4. 实战训练与性能优化

4.1 训练流程配置

python复制def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    
    for img, mask in tqdm(loader):
        img, mask = img.to(device), mask.to(device)
        
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output, mask)
        loss.backward()
        
        # 梯度裁剪
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        total_loss += loss.item()
    
    return total_loss / len(loader)

推荐超参数设置：

4.2 典型问题解决方案

问题1：显存不足

• 解决方案：
  • 使用梯度累积：每4个batch更新一次参数
  • 混合精度训练：torch.cuda.amp.autocast()
  • 减小窗口大小：将默认7×7改为5×5

问题2：边缘分割不连续

• 改进措施：
  • 在损失函数中添加边界加权项
  • 测试时使用滑动窗口重叠预测
  • 后处理使用条件随机场(CRF)

问题3：小器官分割效果差

• 优化方案：
  • 在数据增强中增加小器官的复制粘贴
  • 为不同器官设计类别权重
  • 使用Focal Loss替代标准交叉熵

4.3 模型部署优化

使用TorchScript导出推理模型：

python复制model.eval()
example = torch.rand(1, 1, 224, 224).to(device)
traced_script = torch.jit.trace(model, example)
traced_script.save("swin_unet_medical.pt")

推理性能优化对比：

在临床实际部署中，建议根据硬件条件选择以下方案组合：

• 边缘设备：ONNX量化+OpenVINO
• 云端服务：TensorRT优化
• 研究验证：保持原始PyTorch模型