pytorch RNN:循环神经网络 代码详解 注释详细

import torch
from torch import nn, optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets


# 定义超参数
batch_size = 100
learning_rate = 1e-3
num_epoches = 2

# 下载训练集 MNIST 手写数字训练集
train_dataset = datasets.MNIST(
    root='F:/PycharmProjects/pytorch-beginner-master/02-Logistic Regression/data', train=False, transform=transforms.ToTensor(), download=True)

test_dataset = datasets.MNIST(
    root='F:/PycharmProjects/pytorch-beginner-master/02-Logistic Regression/data', train=False, transform=transforms.ToTensor())

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


# 定义 Recurrent Network 模型
class Rnn(nn.Module):
    def __init__(self, in_dim, hidden_dim, n_layer, n_class):
        super(Rnn, self).__init__()
        self.n_layer = n_layer
        self.hidden_dim = hidden_dim
        #输入[seq_len=100,batch_size=28,input_size=28] 输出为(seq_len=100,batch_size=28,num_directions*hidden_size) hidden_size = hidden_dim=128
        self.lstm = nn.LSTM(in_dim, hidden_dim, n_layer, batch_first=True)
        self.classifier = nn.Linear(hidden_dim, n_class)#输入为[100,128]输出为：[100,10]

    def forward(self, x):
        out, _ = self.lstm(x)#out=[100,28,128]
        out = out[:, -1, :]#out[100,128]
        out = self.classifier(out)#out=[100,10]
        return out


model = Rnn(28, 128, 2, 10)  # 图片大小是28x28
use_gpu = torch.cuda.is_available()  # 判断是否有GPU加速
if use_gpu:
    model = model.cuda()
# 定义loss和optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 开始训练
for epoch in range(num_epoches):
print('epoch {}'.format(epoch + 1))
print('*' * 10)
model.train()
for i, data in enumerate(train_loader, 1):
    img, label = data
    #b, c, h, w = img.size()
    #assert c == 1, 'channel must be 1'
    img = img.squeeze(1)#将输入张量形状中的1 去除 比如输入为[2,1,3,5,1]的5维矩阵，去掉1后变为[2,3,5]的3维矩阵
    # img = img.view(b*h, w)
    # img = torch.transpose(img, 1, 0)
    # img = img.contiguous().view(w, b, -1)
    if use_gpu:
        img = Variable(img).cuda()
        label = Variable(label).cuda()
    else:
        img = Variable(img)
        label = Variable(label)
    # 向前传播
    out = model(img)#输入[seq_len=100,batch_size=28,input_size=28]
    loss = criterion(out, label)

    # 向后传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # if i % 300 == 0:
    #     print("loss: ",  loss.data.item())
print("loss: ", loss.data.item())

model.eval()
for data in test_loader:
    img, label = data
    b, c, h, w = img.size()
    assert c == 1, 'channel must be 1'
    img = img.squeeze(1)
    # img = img.view(b*h, w)
    # img = torch.transpose(img, 1, 0)
    # img = img.contiguous().view(w, b, h)
    if use_gpu:
        img = Variable(img).cuda()
        label = Variable(label).cuda()
    else:
        img = Variable(img)
        label = Variable(label)
    out = model(img)
    loss = criterion(out, label)
print("test loss: ",  loss.data.item())