axial attention

Axial Attention（轴向注意力）是一种注意力机制，通常由行注意力（row-attention）和列注意力（column-attention）组合使用。它在图像处理中被广泛应用。轴向注意力的使用方法如下所示：

from axial_attention import AxialAttention

img = torch.randn(1, 3, 256, 256)
attn = AxialAttention(dim=3)

其中，`dim`参数表示嵌入维度。轴向注意力的主要思想是在图像的垂直和水平方向上分别进行自我注意力计算，这样可以将计算复杂度从O(2*H*W)降低到O(H*W)。

相关问题

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 进行拆分，然后在每个张量上应用一个线性变换。最后，我们将所有张量拼接起来，得到最终的输出张量。

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

Axial Attention in Multidimensional Transformers

Axial Attention in Multidimensional Transformers是一种用于多维Transformer模型的注意力机制。在传统的Transformer模型中，注意力机Axial Attention in Multidimensional Transformers是一种基于轴向注意力的变种transformer。它允许在解码期间并行计算绝大多数上下文，而不引入任何独立假设。这种层结构自然地与编码和解码设置中张量的多维度对齐。Axial Attention in Multidimensional Transformers的用法如下所示：

import torch
from axial_attention import AxialAttention

img = torch.randn(1, 3, 256, 256)
attn = AxialAttention(dim=3, # embedding dimension
                      dim_index=1, # index of the dimension to split and permute
                      heads=8, # number of heads
                      dim_head=None, # dimension of each head, defaults to dim/heads
                      sum_axial_out=True, # whether to sum the axial output
                      sum_axial_out_dims=None, # dimensions to sum the axial output, defaults to all axial dimensions
                      axial_pos_emb=None, # axial position embedding, defaults to None
                      axial_pos_shape=None, # axial position shape, defaults to None
                      axial_pos_emb_dim=None, # axial position embedding dimension, defaults to dim_head
                      attn_drop=0., # attention dropout
                      proj_drop=0.) # projection dropout

相关问题:
1. 什么是transformer？
2. 轴向注意力在哪些领域有应用？