lstm的pytorch代码

下面是一个简单的LSTM模型的PyTorch代码示例：

import torch
import torch.nn as nn

class LSTMModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
        super(LSTMModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

在这个模型中，我们使用了一个单层的LSTM，输入维度为input_dim，输出维度为output_dim。我们还添加了一个全连接层fc用于输出。在forward函数中，我们首先使用h0和c0初始化了LSTM的隐藏状态和记忆状态，然后将输入x传入LSTM，最后使用全连接层fc将LSTM的输出转换为我们所需要的输出。