写一个基于pytorch的LSTM代码

以下是一个基于PyTorch的LSTM代码示例：

import torch
from torch import nn

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTMModel, 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):
        # Initialize hidden state with zeros
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).requires_grad_()
        
        # Initialize cell state with zeros
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).requires_grad_()
        
        # Move to GPU if available
        if torch.cuda.is_available():
            h0 = h0.cuda()
            c0 = c0.cuda()

        # Forward propagate LSTM
        out, _ = self.lstm(x, (h0.detach(), c0.detach()))

        # Decode the hidden state of the last time step
        out = self.fc(out[:, -1, :])
        return out

这个代码定义了一个LSTM模型，输入的大小为`input_size`，隐藏状态的大小为`hidden_size`，LSTM层数为`num_layers`，输出的大小为`output_size`。在`forward`函数中，我们首先初始化了LSTM的隐藏状态和细胞状态，并使用`batch_first=True`来指定输入张量的第一维是batch size，第二维是时间步长。然后我们使用`self.lstm`对输入进行前向传播，并获取最后一个时间步长的输出。最后，我们将最后一个时间步长的输出通过全连接层`self.fc`映射到输出张量的大小。

要使用此模型，您只需执行以下步骤：

# Instantiate the model
model = LSTMModel(input_size=10, hidden_size=20, num_layers=2, output_size=1)

# Move the model to the GPU if available
if torch.cuda.is_available():
    model.cuda()

# Define the loss function and optimizer
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Train the model
for epoch in range(num_epochs):
    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    
    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Print loss every 100 epochs
    if (epoch+1) % 100 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

在这个示例中，我们首先实例化了模型，并将其移动到GPU（如果可用）。然后我们定义了损失函数和优化器。接下来，我们进行训练，每个epoch进行前向传播、计算损失、反向传播和优化。最后，我们打印出每100个epoch的损失。