3D Swin Transformer 入门指南：从零构建高效视觉模型

1次阅读

没有评论

共计 3053 个字符，预计需要花费 8 分钟才能阅读完成。

传统 3D CNN 在处理视频或医学影像时面临两个主要问题：

计算量爆炸 ：3D 卷积核的参数量随输入尺寸立方级增长，例如处理 128x128x128 的 CT 扫描时，单个 3D 卷积层的计算量可能是 2D 情况的数百倍
长程依赖建模困难 ：普通卷积的感受野有限，难以捕捉跨帧或跨切片的全局关系

而 Transformer 架构天生适合建模长序列依赖，但传统 ViT 直接应用于 3D 数据时会出现：

计算复杂度随 token 数量呈平方增长（视频帧数×每帧 patch 数）
完全忽略 3D 数据的局部连续性特征

3D 窗口划分
将输入体积划分为不重叠的局部立方体（如 8x8x8）
每个窗口内独立计算注意力，复杂度从 O(N²) 降到 O(N)
移位窗口注意力
交替使用常规窗口和移位 50% 的窗口
实现跨窗口信息交互而不增加计算量
分层特征提取
通过 patch merging 逐步下采样
形成多尺度特征金字塔，适合密集预测任务

特性	标准 ViT	3D Swin Transformer
计算复杂度	O((THW)²)	O(THW)
位置编码	绝对位置	相对位置偏置
局部性建模	无	窗口内自注意力
适合任务	分类	检测 / 分割

import torch
from einops import rearrange

# 模拟 CT 扫描数据 (B, C, D, H, W)
data = torch.randn(2, 1, 64, 128, 128)  

# 转换为 patch 序列 (B, L, C×P³)
patch_size = 4
data = rearrange(data, 'b c (d p1) (h p2) (w p3) -> b (d h w) (c p1 p2 p3)', 
                p1=patch_size, p2=patch_size, p3=patch_size)

class Swin3DBlock(nn.Module):
    def __init__(self, dim, window_size, shift_size=0):
        super().__init__()
        self.window_size = window_size
        self.shift_size = shift_size

        # 相对位置偏置表 (可学习参数)
        self.relative_position_bias_table = nn.Parameter(torch.zeros((2*window_size-1)**3, num_heads))

        # 移位窗口实现
        if shift_size > 0:
            self.register_buffer("mask", self.create_mask())

    def create_mask(self):
        """生成移位窗口的注意力掩码"""
        D, H, W = self.window_size
        img_mask = torch.zeros((1, D, H, W, 1))

        # 划分移位区域
        slices = [slice(0, -self.window_size),
                 slice(-self.window_size, -self.shift_size),
                 slice(-self.shift_size, None)]

        cnt = 0
        for d in slices:
            for h in slices:
                for w in slices:
                    img_mask[:, d, h, w, :] = cnt
                    cnt += 1

        return rearrange(img_mask, 'b d h w c -> b (d h w) c')

    def forward(self, x):
        B, L, C = x.shape
        D = H = W = int(L ** (1/3))  # 假设是立方体输入

        # 窗口划分
        x = rearrange(x, 'b (d h w) c -> b c d h w', d=D, h=H, w=W)
        if self.shift_size > 0:
            x = torch.roll(x, shifts=(-self.shift_size,)*3, dims=(2,3,4))

        # 计算窗口注意力
        x = rearrange(x, 'b c (d w1) (h w2) (w w3) -> (b d h w) (w1 w2 w3) c',
                     w1=self.window_size, w2=self.window_size, w3=self.window_size)

        # 此处应实现带偏置的注意力计算
        # ...

        # 恢复原始布局
        x = rearrange(x, '(b d h w) (w1 w2 w3) c -> b (d w1) (h w2) (w w3) c',
                     w1=self.window_size, w2=self.window_size, w3=self.window_size,
                     d=D//self.window_size, h=H//self.window_size, w=W//self.window_size)

        if self.shift_size > 0:
            x = torch.roll(x, shifts=(self.shift_size,)*3, dims=(2,3,4))

        return x

梯度检查点

from torch.utils.checkpoint import checkpoint

def forward(self, x):
    x = checkpoint(self.swin_block1, x)  # 不保存中间激活值
    x = self.swin_block2(x)  # 常规方式

混合精度训练

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

使用 warmup 阶段避免早期震荡
对位置偏置参数使用更大学习率（约 10 倍于其他参数）

param_groups = [{"params": [p for n,p in model.named_parameters() 
               if "bias" in n], "lr": lr*10},
    {"params": [p for n,p in model.named_parameters() 
               if "bias" not in n]}
]
optimizer = AdamW(param_groups, lr=5e-5)

Loss 震荡不收敛
检查窗口尺寸是否适合输入分辨率
尝试减小初始学习率（3D 任务通常需要更小的 lr）
显存不足
降低 batch size 至 1 -2

使用梯度累积模拟更大 batch

for i, (inputs, targets) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, targets) / accumulation_steps
    loss.backward()

    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()