软注意力机制复杂代码

以下是一个稍微复杂的实现软注意力机制的代码示例：

import torch
import torch.nn as nn

class SoftAttention(nn.Module):
    def __init__(self, hidden_size, dropout):
        super(SoftAttention, self).__init__()
        self.hidden_size = hidden_size
        self.dropout = dropout
        
        self.attn = nn.Linear(self.hidden_size * 2, self.hidden_size)
        self.v = nn.Linear(self.hidden_size, 1, bias=False)
        self.dropout = nn.Dropout(self.dropout)
        
    def forward(self, encoder_outputs, hidden, mask=None):
        max_len = encoder_outputs.size(0)
        batch_size = encoder_outputs.size(1)
        
        H = hidden.repeat(max_len, 1, 1).transpose(0, 1)
        encoder_outputs = encoder_outputs.transpose(0, 1)
        
        attn_weights = torch.tanh(self.attn(torch.cat([H, encoder_outputs], dim=2)))
        attn_weights = self.v(attn_weights).squeeze(2)
        
        if mask is not None:
            attn_weights = attn_weights.masked_fill(mask == 0, -1e10)
        
        attn_weights = nn.functional.softmax(attn_weights, dim=1)
        attn_weights = self.dropout(attn_weights)
        
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs)
        context = context.transpose(0, 1)
        
        return context, attn_weights

这个代码和之前的代码类似，不过增加了一些细节：

- 增加了一个dropout层，用于在计算注意力权重时进行随机失活，防止过拟合。
- 增加了一个mask参数，用于在计算注意力权重时进行遮盖，防止注意力权重被mask掉的位置影响计算。如果mask参数不为None，则将mask等于0的位置对应的注意力权重设置为一个极小的数，这样在softmax之后这些位置的权重就会接近于0，不会对上下文向量的计算产生影响。
- 在计算上下文向量时，使用了torch.bmm函数进行矩阵乘法，这个函数的输入是两个3D张量，分别是(batch_size, n, m)和(batch_size, m, p)，输出是一个3D张量，大小为(batch_size, n, p)。这里我们将attn_weights和encoder_outputs分别变成了(batch_size, max_len, 1)和(batch_size, max_len, hidden_size)的形状，然后将它们拼接起来得到(batch_size, max_len, hidden_size + 1)的张量，送入线性层计算注意力权重，最终得到大小为(batch_size, max_len, 1)的attn_weights。注意力权重经过softmax和dropout之后，和encoder_outputs做矩阵乘法得到上下文向量context。最后将context和attn_weights返回。