python 循环神经网络RNN实例——手写数字识别

数据集官网：http://yann.lecun.com/exdb/mnist/
参考：https://github.com/MorvanZhou/PyTorch-Tutorial/blob/master/tutorial-contents/402_RNN_classifier.py
代码包含自动下载数据集

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
import os

batch_size = 200    #分批训练数据、每批数据量
learning_rate = 1e-2    #学习率
num_epoches = 10       #训练次数
DOWNLOAD_MNIST = False    #是否网上下载数据

# Mnist digits dataset
if not(os.path.exists('./mnist/')) or not os.listdir('./mnist/'):
    # not mnist dir or mnist is empyt dir
    DOWNLOAD_MNIST = True

train_dataset = datasets.MNIST(
    root = './mnist',
    train= True,        #download train data
    transform = transforms.ToTensor(),
    download=DOWNLOAD_MNIST
)
test_dataset = datasets.MNIST(
    root = './mnist',
    train= False,        #download test data
    transform = transforms.ToTensor(),
    download=DOWNLOAD_MNIST
)

#该接口主要用来将自定义的数据读取接口的输出或者PyTorch已有的数据读取接口的输入
# 按照batch size封装成Tensor，后续只需要再包装成Variable即可作为模型的输入
train_loader = DataLoader(train_dataset, batch_size = batch_size,shuffle=True)    #shuffle 是否打乱加载数据
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False)

class RNNs(nn.Module):
    def __init__(self,in_dim,hidden_dim,n_layer,n_class):
        super(RNNs,self).__init__()
        self.n_layer = n_layer   #RNNs 网络的层数，不包括全连接层
        self.hidden_dim = hidden_dim    #RNNs网络的输出维度，也是全连接层的输入维度

        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,hidden = self.LSTM(x)   #Outputs: output, (h_n, c_n)
        out = out[:,-1,:]      #一张图片有28个相关联的输入，只需要最后一个输入的输出状态
        out = self.classifier(out)
        return out

model = RNNs(28,128,2,10)   #每张图片(28*28)被看作28个连续的输入，每个输入是28维（即in_dim)的向量
if torch.cuda.is_available():       #是否可用GPU计算
     model = model.cuda()           #转换成可用GPU计算的模型

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
#常用优化方法有
#1.Stochastic Gradient Descent (SGD)
#2.Momentum
#3.AdaGrad
#4.RMSProp
#5.Adam (momentum+adaGrad)   效果较好
for epoch in range(num_epoches):
    print('epoch {}'.format(epoch+1))
    print('*'*10)
    running_loss = 0.0
    running_acc = 0.0
    #训练
    for i,data in enumerate(train_loader,1):
        img,label = data
        img = img.squeeze(1)  #去除size为1的维度，比如A.size为(4*1*3)，则A.squeeze(1)的size为（4*3）

        #判断是否可以使用GPU，若可以则将数据转化为GPU可以处理的格式。
        if torch.cuda.is_available():
            img = Variable(img).cuda()
            label = Variable(label).cuda()
        else:
            img = Variable(img)
            label = Variable(label)

        out = model(img)
        loss = criterion(out,label)
        running_loss += loss.data.item() * label.size(0)
        _,pred = torch.max(out,1)
        num_correct = (pred == label).sum()
        running_acc += num_correct.data.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i % 300 == 0:
            print('[{}/{}] Loss: {:.6f}, Acc: {:.6f}'.format(
                epoch + 1, num_epoches, running_loss / (batch_size * i),
                running_acc / (batch_size * i)))

    print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format( \
        epoch + 1, running_loss / len(train_dataset), running_acc / len(train_dataset)))

##省略测试代码，与训练代码相似。