Pytorch学习之LSTM识别MNIST数据集

引入库函数
引入pytorch库，主要是nn,optim,Variable。

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

设置超参数
设置超参数，其为神经网络的基础配置。

# 超参数
batch_size = 100    # 批大小
learning_rate = 0.01 # 学习率
num_epoches = 20  # 训练次数

数据集预处理
将图片转换成tensor类型，并把图片中心化，缩放到[-1,1]上。

data_tf = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5],[0.5])])

获得数据
下载MINIST集数据，并获取训练和测试集数据。

train_dataset = datasets.MNIST(root='./data',train=True,transform=data_tf,download=True)
test_dataset = datasets.MNIST(root='./data',train=False,transform=data_tf)
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False)

建立RNN模型
建立RNN神经网络，该神经网络为两层，第一层为LSTM神经网络，用于进行序列分析，第二层为全连接层，用作分类。

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
        self.lstm = nn.LSTM(in_dim,hidden_dim,n_layer,batch_first=True)
        self.classifier = nn.Linear(hidden_dim,n_class)
    
    def forward(self,x):
        out,_ = self.lstm(x)
        out = out[:,-1,:]
        out = self.classifier(out)
        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)

训练数据和评估
训练num_epoches次，每次记录损失值及进行模型测试

# 开始训练
for epoch in range(num_epoches):
    running_loss = 0.0
    running_acc = 0.0
    for i, data in enumerate(train_loader, 1):
        img, label = data
        img = img.squeeze(1)
        if torch.cuda.is_available():
            img = img.cuda()
            label = label.cuda()
        else:
            img = Variable(img)
            label = Variable(label)

            # 向前传播
            out = model(img)
            loss = criterion(out, label)
            running_loss += loss.data[0] * label.size(0)
            _, pred = torch.max(out, 1)
            num_correct = (pred == label).sum()
            running_acc += num_correct.data[0]
            
            # 向后传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if i % 300 == 0:
                print('[{}/{}] Loss: {:.6f}, Acc: {:.6f}'.format(
                    epoch + 1, num_epoches, running_loss.item() / (batch_size * i),
                    running_acc.item() / (batch_size * i)))
        print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(
            epoch + 1, running_loss.item() / (len(train_dataset)), running_acc.item() / (len(
                train_dataset))))
                
    model.eval()
    eval_loss = 0.
    eval_acc = 0.
    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, volatile=True).cuda()
            label = Variable(label, volatile=True).cuda()
        else:
            img = Variable(img, volatile=True)
            label = Variable(label, volatile=True)
        out = model(img)
        loss = criterion(out, label)
        eval_loss += loss.data[0] * label.size(0)
        _, pred = torch.max(out, 1)
        num_correct = (pred == label).sum()
        eval_acc += num_correct.data[0]
    print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss.item() / (len(
        test_dataset)), eval_acc.item() / (len(test_dataset))))

模型结果

[1/20] Loss: 0.276766, Acc: 0.917283  
Test Loss: 0.227039, Acc: 0.931700
[2/20] Loss: 0.368415, Acc: 0.889717
Finish 2 epoch, Loss: 0.368415, Acc: 0.889717
Test Loss: 0.401789, Acc: 0.882400
[3/20] Loss: 0.308447, Acc: 0.907267
Finish 3 epoch, Loss: 0.308447, Acc: 0.907267
Test Loss: 0.254289, Acc: 0.923200
[4/20] Loss: 0.310579, Acc: 0.905850
Finish 4 epoch, Loss: 0.310579, Acc: 0.905850
Test Loss: 0.308438, Acc: 0.904400
[5/20] Loss: 0.400899, Acc: 0.878133
Finish 5 epoch, Loss: 0.400899, Acc: 0.878133
Test Loss: 0.391112, Acc: 0.883500

...

[11/20] Loss: 0.232482, Acc: 0.928167
Finish 11 epoch, Loss: 0.232482, Acc: 0.928167
Test Loss: 0.179359, Acc: 0.944000

...

[13/20] Loss: 0.209641, Acc: 0.934883
Finish 13 epoch, Loss: 0.209641, Acc: 0.934883
Test Loss: 0.187486, Acc: 0.943700

...

[19/20] Loss: 0.304291, Acc: 0.906167
Finish 19 epoch, Loss: 0.304291, Acc: 0.906167
Test Loss: 0.303386, Acc: 0.910300
[20/20] Loss: 0.312536, Acc: 0.903133
Finish 20 epoch, Loss: 0.312536, Acc: 0.903133
Test Loss: 0.349427, Acc: 0.885400

更多计算机视觉与图形学相关资料，请关注微信公众号：计算机视觉与图形学实战