3D Transformer入门指南：从零构建你的第一个空间注意力模型

1次阅读

没有评论

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

传统 2D Transformer 在图像领域表现出色，但面对 CT/MRI 医疗影像、LiDAR 点云等三维数据时，直接平铺 3D 数据会丢失空间结构信息。3D Transformer 通过以下方式解决这个问题：

保留空间关系 ：体素(voxel) 级别的处理维持了切片间的解剖结构
跨模态融合：可同时处理 PET/CT 等多通道体积数据
长程依赖建模：相比 CNN 的局部感受野，能捕捉病灶间的全局关联

2D 位置编码公式：
$$PE_{(x,y)}^{2D}=\sin\left(\frac{x}{10000^{2i/d}}\right) + \cos\left(\frac{y}{10000^{2i/d}}\right)$$

扩展到 3D 需增加 z 轴分量：
$$PE_{(x,y,z)}^{3D}=PE_{(x,y)}^{2D} + \sin\left(\frac{z}{10000^{2i/d}}\right)$$

假设输入尺寸为 $H \times W$（2D）和 $H \times W \times D$（3D）：

2D 自注意力复杂度：$O(H^2W^2)$
3D 自注意力复杂度：$O(H^2W^2D^2)$

这意味着 16×16×16 的 3D 输入计算量已接近 256×256 的 2D 输入

import torch
import torch.nn as nn

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)
        x = x.flatten(2).transpose(1, 2)  # [B, num_patches, embed_dim]
        return x

关键点说明：

使用 3D 卷积实现体素块切分
典型参数：BraTS 数据集常用 16×16×16 的 patch 大小
输出维度为 $\frac{D}{patch} \times \frac{H}{patch} \times \frac{W}{patch}$

from torch.utils.checkpoint import checkpoint

class MemoryEfficientAttention(nn.Module):
    def forward(self, x):
        # 启用梯度检查点
        return checkpoint(self._attention, x) 

    def _attention(self, x):
        # 实际注意力计算逻辑
        qkv = self.qkv(x).chunk(3, dim=-1)
        attn = (qkv[0] @ qkv[1].transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        return attn @ qkv[2]

Batch Size	输入尺寸	显存占用
1	128³	6.2GB
2	128³	11.8GB
4	128³	OOM

模型	Dice Score	HD95(mm)
3D U-Net	0.891	8.7
3D Transformer	0.902	6.3

错误做法：

# 直接对原始坐标除最大值
coord = coord / 255.0  # 导致不同病例尺度不一致

正确做法：

# 使用各向同性归一化
coord = (coord - coord.mean()) / (coord.std() + 1e-8)

推荐使用 DistributedDataParallel 配合：

torch.distributed.init_process_group(backend='nccl')
model = DDP(model, device_ids=[local_rank])
# 数据加载需配合 sampler
sampler = DistributedSampler(dataset)