3D Transformer 技术解析：从基础原理到高效实现

1次阅读

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

3D Transformer 是 Transformer 架构在三维数据领域的扩展，它能够处理如点云、医学影像（CT/MRI）、3D 视频等数据。与传统的 2D Transformer 相比，3D Transformer 引入了额外的空间维度，使其能够捕捉更丰富的空间关系。

核心特点：
处理三维数据（如体素网格或点云）
通过自注意力机制建模长距离依赖关系
适用于需要全局上下文理解的 3D 任务
典型应用：
医学图像分割（如器官或肿瘤定位）
3D 目标检测（自动驾驶场景理解）
点云补全与分类
3D 生成模型（如 3D-GAN）

2D Transformer 主要处理图像或序列数据，而 3D Transformer 需要额外考虑深度维度：

数据表示：
2D：[batch, height, width, channels]
3D：[batch, depth, height, width, channels]
计算复杂度：
2D 注意力复杂度：O(H²W²C)
3D 注意力复杂度：O(D²H²W²C)（需优化策略）
位置编码：
3D 需要扩展为三维坐标编码
示例：将 2D 的 (x,y) 正弦编码扩展为(x,y,z)

3D 自注意力的关键实现步骤：

输入处理：
将 5D 输入张量展平为[batch, d*h*w, channels]
添加可学习的三维位置编码
注意力计算：
通过线性层生成 Q /K/V
计算缩放点积注意力：softmax(QK^T/√d)V
内存优化技巧：
使用轴向注意力（分 x /y/ z 轴计算）
窗口化注意力（限制局部感受野）

import torch
import torch.nn as nn
import math

class PositionalEncoding3D(nn.Module):
    """三维位置编码"""
    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        inv_freq = 1. / (10000 ** (torch.arange(0, channels, 2).float() / channels))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x):
        # x: [B, D, H, W, C]
        b, d, h, w, c = x.shape
        pos_z = torch.arange(d, device=x.device).type_as(self.inv_freq)
        pos_y = torch.arange(h, device=x.device).type_as(self.inv_freq)
        pos_x = torch.arange(w, device=x.device).type_as(self.inv_freq)

        sin_z = torch.sin(pos_z[:, None] * self.inv_freq[None, :])
        cos_z = torch.cos(pos_z[:, None] * self.inv_freq[None, :])
        sin_y = torch.sin(pos_y[:, None] * self.inv_freq[None, :])
        cos_y = torch.cos(pos_y[:, None] * self.inv_freq[None, :])
        sin_x = torch.sin(pos_x[:, None] * self.inv_freq[None, :])
        cos_x = torch.cos(pos_x[:, None] * self.inv_freq[None, :])

        pos_enc = torch.zeros((d, h, w, c), device=x.device)
        pos_enc[..., 0::2] = (sin_z[:, None, None] + sin_y[:, None] + sin_x).unsqueeze(-1)
        pos_enc[..., 1::2] = (cos_z[:, None, None] + cos_y[:, None] + cos_x).unsqueeze(-1)
        return x + pos_enc.unsqueeze(0)

class Attention3D(nn.Module):
    """3D 自注意力模块"""
    def __init__(self, dim, heads=8, dropout=0.):
        super().__init__()
        self.heads = heads
        self.scale = (dim // heads) ** -0.5

        self.to_qkv = nn.Linear(dim, dim * 3)
        self.to_out = nn.Sequential(nn.Linear(dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        # x: [B, D*H*W, C]
        b, n, c = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: t.view(b, n, self.heads, c // self.heads).transpose(1, 2), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
        attn = dots.softmax(dim=-1)

        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).reshape(b, n, c)
        return self.to_out(out)