门控轴向自注意力机制代码

你可以使用以下代码实现门控轴向自注意力机制（Gated Axial Self-Attention）：

import torch
import torch.nn as nn

class GatedAxialSelfAttention(nn.Module):
    def __init__(self, dim, max_length):
        super(GatedAxialSelfAttention, self).__init__()

        self.dim = dim
        self.max_length = max_length

        # 定义参数矩阵
        self.W_q = nn.Linear(dim, dim, bias=False)
        self.W_k = nn.Linear(dim, dim, bias=False)
        self.W_v = nn.Linear(dim, dim, bias=False)
        self.W_g = nn.Linear(dim, dim, bias=False)
        self.W_o = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        batch_size, seq_length, _ = x.size()

        q = self.W_q(x)  # (batch_size, seq_length, dim)
        k = self.W_k(x)  # (batch_size, seq_length, dim)
        v = self.W_v(x)  # (batch_size, seq_length, dim)

        # 水平方向的自注意力
        attention_h = torch.matmul(q.transpose(1, 2), k.transpose(1, 2).contiguous())  # (batch_size, dim, seq_length)
        attention_h = attention_h.softmax(dim=-1)

        # 竖直方向的自注意力
        attention_v = torch.matmul(k.transpose(1, 2).contiguous(), q.transpose(1, 2))  # (batch_size, seq_length, seq_length)
        attention_v = attention_v.softmax(dim=-1)

        # 门控机制
        g_h = self.W_g(attention_h.transpose(1, 2).contiguous())  # (batch_size, seq_length, dim)
        g_v = self.W_g(attention_v)  # (batch_size, seq_length, dim)

        # 融合水平和竖直的注意力
        combined_attention = attention_h.transpose(1, 2).contiguous() * g_h + attention_v * g_v

        # 输出
        output = torch.matmul(combined_attention, v.transpose(1, 2))  # (batch_size, seq_length, dim)
        output = self.W_o(output)

        return output

# 使用示例
dim = 256
max_length = 100
batch_size = 32
seq_length = 50

input_data = torch.randn(batch_size, seq_length, dim)
attention = GatedAxialSelfAttention(dim, max_length)
output = attention(input_data)
print(output.size())

这个代码实现了一个门控轴向自注意力机制的模块，输入是一个三维张量 `(batch_size, seq_length, dim)`，其中 `batch_size` 是批量大小，`seq_length` 是序列长度，`dim` 是特征维度。模块将输入通过门控轴向自注意力机制进行处理，并返回输出 `(batch_size, seq_length, dim)`。