3D自注意力机制入门指南：从理论到PyTorch实战

1次阅读

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

传统 2D 注意力在处理视频（时序图像）、CT 扫描（空间体数据）时会遇到两个致命问题：

维度缺失：2D 卷积核无法捕捉帧与帧之间的时间关联，比如视频中人物的连续动作
计算浪费：对 3D 数据展开成 2D 处理时（如把视频帧拼成长图），会破坏原始空间拓扑关系

举个实际例子：在肺癌检测任务中，2D 注意力只能单独分析每个 CT 切片，而 3D 注意力可以同时观察肿瘤在 xyz 三个方向的生长趋势。

公式看起来复杂，其实就是在三个维度上重复 2D 注意力的计算逻辑：

Attention(Q,K,V) = softmax(QK^T/√d) V

具体实现时要注意张量维度的变化（以 batch_size= 8 的 16 帧 224×224 视频为例）：

输入数据：[8, 3, 16, 224, 224]（batch, channel, depth, height, width）
线性变换 后得到 Q /K/V：[8, 16×224×224, 64]（展平空间维度）
分头计算：[8, 4, 16×224×224, 16]（4 个注意力头）
输出还原：[8, 3, 16, 224, 224]（保持输入维度）

（注：此处应为示意图描述）

关键实现技巧都封装在这个可配置的模块中：

class Attention3D(nn.Module):
    def __init__(self, dim, heads=4, chunk_size=32):
        super().__init__()
        self.heads = heads
        self.chunk_size = chunk_size  # 内存优化关键参数

        # 线性变换层
        self.to_qkv = nn.Linear(dim, dim * 3) 
        self.to_out = nn.Linear(dim, dim)

    def forward(self, x):
        b, c, d, h, w = x.shape
        x = x.flatten(2).transpose(1, 2)  # [b, d*h*w, c]

        # 分块处理避免 OOM
        if d*h*w > self.chunk_size**3:
            return self.chunked_forward(x)

        # 常规注意力计算
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * (q.shape[-1] ** -0.5)
        attn = dots.softmax(dim=-1)
        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out).view(b, c, d, h, w)

在 RTX 3090 上的测试数据（单位毫秒）：