pytorch实现LSTM

下面是一个简单的 PyTorch LSTM 的实现示例：

import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device=x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device=x.device)

        out, (hn, cn) = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

在上面的代码中，我们定义了一个 `LSTM` 类，它继承了 `nn.Module` 类。该类包含了 LSTM 层和全连接层。在 `__init__` 中，我们定义了 LSTM 层和全连接层的参数。在 `forward` 中，我们先定义了初始的隐藏状态和细胞状态为全零，并将其发送到 LSTM 层中。最后，我们将 LSTM 层的输出传递到全连接层中，并返回输出。

可以使用以下代码来测试这个 LSTM 模型：

input_size = 10
hidden_size = 20
num_layers = 2
output_size = 1

batch_size = 3
seq_len = 5

x = torch.randn(batch_size, seq_len, input_size)
lstm = LSTM(input_size, hidden_size, num_layers, output_size)

print(lstm(x))

在上面的代码中，我们生成了一个大小为 `(batch_size, seq_len, input_size)` 的张量 `x`，将其传递到 LSTM 模型中，并打印输出结果。