3D Vision Transformer 实战：PyTorch 代码实现与新手避坑指南

1次阅读

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

在 3D 视觉任务（如医学影像分析、点云处理）中，传统 CNN 面临三大局限：

感受野受限：卷积核难以捕捉长距离依赖关系
参数效率低：3D 卷积核导致参数量立方级增长
各向异性处理：对空间不同方向的特征提取能力不均

Vision Transformer 通过以下优势解决这些问题：

全局注意力机制自然建模体素间关系
权重共享降低参数量
位置编码保持空间感知

组件	2D 版本	3D 版本
Patch Embedding	(H,W)→(N,P²×C)	(D,H,W)→(N,P³×C)
Position Encoding	2D 正弦位置编码	3D 正弦位置编码
Attention 计算	(B,N,C)矩阵运算	(B,N,C)需处理更高维度

class PatchEmbed3D(nn.Module):
    def __init__(self, img_size=128, patch_size=16, in_chans=1, embed_dim=768):
        super().__init__()
        self.proj = nn.Conv3d(in_chans, embed_dim, 
                            kernel_size=patch_size, 
                            stride=patch_size)

    def forward(self, x):
        # x: (B, C, D, H, W)
        x = self.proj(x)  # (B, E, D/p, H/p, W/p)
        x = x.flatten(2).transpose(1, 2)  # (B, N, E)
        return x

位置编码公式：

$$
PE_{(d,h,w,2i)} = sin(\frac{d}{10000^{2i/D}}) + sin(\frac{h}{10000^{2i/D}}) + sin(\frac{w}{10000^{2i/D}})
$$

实现代码：

def get_3d_position_embedding(grid_size, dim):
    depth, height, width = grid_size
    position_ids = torch.stack(torch.meshgrid(torch.arange(depth), 
        torch.arange(height),
        torch.arange(width)
    ), dim=-1).float()

    div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))
    pe = torch.zeros(depth, height, width, dim)
    pe[..., 0::2] = torch.sin(position_ids @ div_term)
    pe[..., 1::2] = torch.cos(position_ids @ div_term)
    return pe.view(-1, dim)  # (D*H*W, C)

关键实现要点：

使用 einops 库高效处理维度变换
注意力分数计算时需 mask 无效体素
采用 pre-norm 结构增强训练稳定性

import einops

class MultiHeadAttention3D(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.scale = (dim // num_heads) ** -0.5

        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = einops.rearrange(self.qkv(x), 
            "b n (qkv h c) -> qkv b h n c", 
            qkv=3, h=self.num_heads
        )

        attn = (qkv[0] @ qkv[1].transpose(-2, -1)) * self.scale
        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)
        attn = attn.softmax(dim=-1)

        x = (attn @ qkv[2])
        x = einops.rearrange(x, "b h n c -> b n (h c)")
        return self.proj(x)

推荐使用 torchio 进行 3D 医学影像处理：

import torchio as tio

transforms = tio.Compose([tio.RandomFlip(axes=(0, 1, 2)),
    tio.RandomAffine(scales=(0.9, 1.1)),
    tio.RandomBlur(),
    tio.ZNormalization()])

dataset = tio.SubjectsDataset(
    subjects_list,
    transform=transforms
)

完整模型架构示例：

class VisionTransformer3D(nn.Module):
    def __init__(self, img_size=128, patch_size=16, in_chans=1, num_classes=10):
        super().__init__()
        self.patch_embed = PatchEmbed3D(img_size, patch_size, in_chans)
        grid_size = (img_size // patch_size,) * 3
        self.pos_embed = nn.Parameter(get_3d_position_embedding(grid_size, 768))

        self.blocks = nn.ModuleList([nn.TransformerEncoderLayer(d_model=768, nhead=12) 
            for _ in range(12)
        ])

        self.head = nn.Linear(768, num_classes)

    def forward(self, x):
        x = self.patch_embed(x)
        x = x + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        return self.head(x.mean(1))

使用 torch.utils.checkpoint 进行梯度检查点
调整 batch_size 时保持 D*H*W 乘积恒定
启用 pin_memory 加速数据加载

scaler = torch.cuda.amp.GradScaler()

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

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

典型错误场景：

输入维度错误：
预期：(B,C,D,H,W)
实际：(B,D,H,W,C)
修复：x = x.permute(0,4,1,2,3)
位置编码维度不符：
检查 grid_size 计算是否正确

推荐采用分层学习率：

param_groups = [{"params": model.patch_embed.parameters(), "lr": base_lr},
    {"params": model.pos_embed, "lr": base_lr * 0.1},
    {"params": model.head.parameters(), "lr": base_lr * 5}
]
optimizer = AdamW(param_groups)

在医疗影像领域的潜在应用方向：