LSTM+attention机制 Python代码 Pytorch框架

以下是LSTM attention机制的Python代码（使用Pytorch框架）：

import torch
import torch.nn as nn
import torch.nn.functional as F

class LSTMAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LSTMAttention, self).__init__()
        self.hidden_dim = hidden_dim
        self.lstm = nn.LSTM(input_dim, hidden_dim, bidirectional=True)
        self.fc = nn.Linear(hidden_dim * 2, output_dim)
        self.attention_fc = nn.Linear(hidden_dim * 2, 1)

    def forward(self, x):
        # Initialize hidden state and cell state
        hidden_state = torch.zeros(2, x.size(0), self.hidden_dim).to(x.device)
        cell_state = torch.zeros(2, x.size(0), self.hidden_dim).to(x.device)

        lstm_out, _ = self.lstm(x, (hidden_state, cell_state))

        # Apply attention
        attention_weights = F.softmax(self.attention_fc(lstm_out), dim=1)
        attention_vector = torch.sum(attention_weights * lstm_out, dim=1)

        # Feed attention_vector to the fully connected layer
        out = self.fc(attention_vector)

        return out

该代码定义了一个带有attention机制的LSTM模型，输入维度为input_dim，隐层维度为hidden_dim，输出维度为output_dim。模型首先通过LSTM层对数据进行处理，然后通过注意力机制计算加权平均值，最后将得到的注意力向量馈送给全连接层以生成模型的输出。