基于pytorch的rnn_lstm模型代码

好的，下面是一个基于PyTorch框架的RNN-LSTM模型的代码示例，用于实现MNIST手写数字分类：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader

# 准备数据集
train_dataset = MNIST(root='./data', train=True, transform=ToTensor(), download=True)
test_dataset = MNIST(root='./data', train=False, transform=ToTensor())

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

# 构建模型
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(RNN, 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, num_classes)

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

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

model = RNN(28, 128, 2, 10).cuda()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = Variable(images.view(-1, 28, 28)).cuda()
        labels = Variable(labels).cuda()

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f' % (epoch+1, num_epochs, i+1, len(train_dataset)//128, loss.data[0]))

# 评估模型
correct = 0
total = 0
for images, labels in test_loader:
    images = Variable(images.view(-1, 28, 28)).cuda()
    outputs = model(images)
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted.cpu() == labels).sum()

print('Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total))

该模型包含一个LSTM层和一个全连接层，其中LSTM层的输入形状为(28,28)，表示每个手写数字图像的像素为28x28。在模型训练时，使用交叉熵作为损失函数，Adam作为优化器。在训练过程中，使用128个样本的批量训练，共进行10个epoch的训练。最后输出测试集上的准确率。需要注意的是，该代码示例使用了GPU进行训练，如果没有GPU可以将`.cuda()`去掉。