LSTM代码 pytorch

LSTM是一种循环神经网络，它可以处理序列数据，并且在处理长序列时比传统的循环神经网络更有效。以下是基于pytorch实现的LSTM代码：

import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, batch_first=True):
        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=batch_first)

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

这个代码实现了一个单向的LSTM模型，其中包含了输入大小，隐藏层大小和层数等参数。在forward函数中，我们首先初始化了LSTM的初始隐藏状态和细胞状态，然后将输入x传递给LSTM模型，最后返回输出out。

如果需要实现双向LSTM，只需要将nn.LSTM替换为nn.LSTMBiDirectional即可。