axial attention代码详述

Axial Attention 是一种用于处理序列数据的注意力机制，它在注意力计算时将序列数据分解成多个轴，然后在每个轴上进行注意力计算。这种方法可以有效地降低注意力计算的复杂度，同时在处理长序列数据时也能够取得较好的效果。

在 PyTorch 中，可以通过自定义 nn.Module 来实现 Axial Attention。以下是一个简单的实现示例：

import torch
import torch.nn as nn

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

        self.axial_dims = axial_dims
        self.axial_linears = nn.ModuleList([
            nn.Linear(dim, dim) for _ in axial_dims
        ])

        self.attention_linear = nn.Linear(dim, dim)
        self.out_linear = nn.Linear(dim, dim)

    def forward(self, x):
        # x: [batch_size, seq_len, dim]

        # split into multiple axial tensors
        tensors = []
        for i, dim in enumerate(self.axial_dims):
            tensors.append(x.transpose(1, i + 2).contiguous().view(-1, dim, self.head_dim))
        tensors = torch.cat(tensors, dim=1)

        # calculate attention scores
        attn_scores = self.attention_linear(tensors)
        attn_scores = attn_scores.transpose(1, 2).contiguous()
        attn_scores = attn_scores.view(-1, self.num_heads, tensors.shape[1], tensors.shape[2])
        attn_scores = attn_scores.softmax(dim=-1)
        attn_scores = attn_scores.flatten(start_dim=2, end_dim=3)

        # calculate weighted sum of values
        values = tensors.view(-1, self.num_heads, tensors.shape[1], tensors.shape[2])
        attn_output = torch.einsum('bhnd,bhdf->bhnf', attn_scores, values)
        attn_output = attn_output.flatten(start_dim=2)

        # merge axial tensors and apply linear transformations
        out = self.out_linear(attn_output)
        for i, dim in enumerate(self.axial_dims):
            out = out.view(-1, dim, self.head_dim)
            out = self.axial_linears[i](out)
        out = out.transpose(1, len(self.axial_dims) + 1)

        return out

在这个实现中，我们首先将输入张量按照 axial_dims 进行分解，然后在每个分解出来的张量上进行注意力计算。在计算注意力时，我们首先利用一个线性变换将张量中的每个元素映射到注意力空间中，然后计算每个元素与其它元素的相似性得分。这里我们使用了 softmax 函数来将得分归一化到 [0, 1] 的范围内。最后，我们将得分与张量中的每个元素进行加权求和，得到注意力输出。

注意力输出张量的形状与输入张量相同，因此我们需要将注意力输出张量重新合并成一个张量。在这个实现中，我们首先将注意力输出张量按照 axial_dims 进行拆分，然后在每个张量上应用一个线性变换。最后，我们将所有张量拼接起来，得到最终的输出张量。

需要注意的是，这个实现中并没有包含位置编码等常用的序列建模技巧，因此在实际应用中需要根据具体情况进行调整。