用python写一个基于pytorch框架的lstm模型

好的，以下是一个基于PyTorch框架的LSTM模型的示例代码：

import torch
import torch.nn as nn

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(input_size, hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
        
    def forward(self, input):
        batch_size = input.size(0)
        hidden = self.init_hidden(batch_size)
        output, hidden = self.lstm(input, hidden)
        output = self.fc(output[:, -1, :])
        return output
        
    def init_hidden(self, batch_size):
        hidden = (torch.zeros(1, batch_size, self.hidden_size),
                  torch.zeros(1, batch_size, self.hidden_size))
        return hidden

在这个模型中，我们定义了一个LSTM模型类，包含了一个LSTM层和一个全连接层。在forward方法中，我们首先初始化了LSTM的隐藏状态，然后将输入输入到LSTM中，得到最后一个时间步的输出。最后，我们将最后一个时间步的输出传入全连接层中，得到最终的输出。

需要注意的是，我们在init_hidden方法中初始化了LSTM的隐藏状态，这个方法会返回一个元组，包含了LSTM的两个隐藏状态。在每次调用LSTM的forward方法时，我们需要将这个隐藏状态传入，这样LSTM才能保留之前的状态信息。

以上是一个简单的LSTM模型的实现，你可以根据自己的需要进行调整和扩展。