3D局部窗口自注意力机制入门指南：从原理到PyTorch实现

1次阅读

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

在 3D 视觉任务中，全局自注意力面临立方级复杂度增长。假设输入体积尺寸为 D×H×W，计算复杂度为：

$$O((D \times H \times W)^2)$$

例如处理 128×128×128 的 CT 扫描时，单层注意力需要约 274 亿次运算，显存占用超过 40GB。这导致：

无法处理高分辨率数据
训练 batch_size 被压缩至 1 -2
长序列梯度不稳定

类型	FLOPs 示例(64³输入)	内存占用	感受野
全局注意力	17.2TFLOPS	256GB	全图
滑动窗口(8³)	0.8TFLOPS	6GB	局部 8×8×8
膨胀窗口(4 倍)	1.2TFLOPS	9GB	32×32×32
轴向注意力	0.3TFLOPS	3GB	全局 + 局部混合

import torch
import einops

def window_partition(x, window_size):
    """
    Args:
        x: (B, C, D, H, W)
        window_size: (wd, wh, ww)
    Returns:
        windows: (B*num_windows, window_size*window_size*window_size, C)
    """
    B, C, D, H, W = x.shape
    x = x.view(B, C, 
               D // window_size[0], window_size[0], 
               H // window_size[1], window_size[1],
               W // window_size[2], window_size[2])
    x = einops.rearrange(x, 
                        'b c d1 wd h1 wh w1 ww -> (b d1 h1 w1) (wd wh ww) c')
    return x

def get_relative_positions(window_size):
    coords = torch.stack(torch.meshgrid([torch.arange(window_size[i]) for i in range(3)]), dim=-1)
    relative_coords = coords[:, None] - coords[None, :]  # (wh*ww*wd, wh*ww*wd, 3)
    return relative_coords + window_size[0] - 1  # 转换为正数

class RelativePositionBias(nn.Module):
    def __init__(self, window_size, num_heads):
        super().__init__()
        self.relative_position_bias_table = nn.Parameter(torch.zeros((2 * window_size[0] - 1) * 
                       (2 * window_size[1] - 1) * 
                       (2 * window_size[2] - 1), num_heads))
        self.register_buffer('relative_positions', 
                            get_relative_positions(window_size).view(-1,3))

    def forward(self):
        relative_pos_index = self.relative_positions @ torch.tensor([(2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1), 
             (2 * self.window_size[2] - 1), 1])
        return self.relative_position_bias_table[relative_pos_index]

def optimize_memory_format(model, input_tensor):
    model = model.to(memory_format=torch.channels_last_3d)
    input_tensor = input_tensor.contiguous(memory_format=torch.channels_last_3d)
    return model, input_tensor  # 提升 30% 推理速度

from torch.utils.checkpoint import checkpoint

def forward_with_checkpoint(self, x):
    def create_custom_forward(module):
        def custom_forward(*inputs):
            return module(inputs[0])
        return custom_forward

    return checkpoint(create_custom_forward(self.attention_block), 
                     x, use_reentrant=False)

当窗口大小不能被特征图尺寸整除时：

使用 nn.ConstantPad3d 进行对称填充
在反向传播前用 x = x[:, :, :D, :H, :W] 裁剪回原尺寸

def adaptive_padding(x, window_size):
    pad_d = (window_size[0] - x.size(2) % window_size[0]) % window_size[0]
    pad_h = (window_size[1] - x.size(3) % window_size[1]) % window_size[1]
    pad_w = (window_size[2] - x.size(4) % window_size[2]) % window_size[2]
    return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_d))

class WindowAttention3D(nn.Module):
    def __init__(self, dim, window_size, num_heads):
        super().__init__()
        self.dim = dim
        self.window_size = window_size
        self.num_heads = num_heads

        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)
        self.relative_position_bias = RelativePositionBias(window_size, num_heads)

    def forward(self, x):
        B_, N, C = x.shape  # N = wd*wh*ww
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # [B_, num_heads, N, C/num_heads]

        attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        attn = attn + self.relative_position_bias()

        attn = attn.softmax(dim=-1)
        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        return self.proj(x)