Pyorch 快速搭建手写识别网络 LeNet5

Pytorch 快速搭建手写识别网络

Pytorch是一个十分简洁的神经网络框架，可以快速的构建出需要的机器学习网络。

Pytorch的安装和简介请看：https://blog.csdn.net/qq_33302004/article/details/106320649

本文以手写识别问题为基础，搭建LeNet 5网络，网络结构如下图：

与图中不同的是我的数据集中输入数据为28*28（而不是32*32），所以第一个全连接层为16*4*4——120。代码如下：

#coding=utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision                      # 包含有支持加载类似Imagenet，CIFAR10，MNIST 等公共数据集的数据加载模块 
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np

# 超参数
batch_size = 100          # 批数量
learning_rate = 0.01     # 学习率
momentum = 0.5           # 冲量
TRAINING_STEPS = 10      # 训练次数？


# 数据预处理
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(0.5,0.5,0.5)])

# 读取数据
dataset_train = datasets.MNIST('MNIST/',train=True,transform = transforms.ToTensor())
dataset_test = datasets.MNIST('MNIST/',train=False,transform = transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(
    dataset = dataset_train,
    batch_size = batch_size,
    shuffle = True              # 将训练模型的数据集进行打乱的操作
)
test_loader = torch.utils.data.DataLoader(
    dataset = dataset_test,
    batch_size = batch_size,
    shuffle = False
)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(1,6,5)       # 第一层： 输入通道数为1，输出通道数为6，卷积核5*5
        self.pool1 = nn.MaxPool2d(2)        # 第二层： 最大池化
        self.conv2 = nn.Conv2d(6,16,5)      # 第三层： 输入通道数为6，输出通道数为16，卷积核5*5
        self.pool2 = nn.MaxPool2d(2)        # 第四层： 最大池化
        self.fc1 = nn.Linear(16*4*4,120)    # 第五层： 全连接，输入节点16*5*5，输出节点120
        self.fc2 = nn.Linear(120,84)        # 第六层： 全连接，输入节点120，输出节点84
        self.fc3 = nn.Linear(84,10)         # 第七层： 全连接，输入节点84，输出节点10

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(-1, 16*4*4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        x = F.log_softmax(x,dim = 1)
        return x

def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 网络实例化
net = Net()
# 构建优化器
optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=momentum)  # weight decay 什么意思

def do_train(epoch):
    net.train()
    for step, (x, y) in enumerate(train_loader):
        x, y = Variable(x), Variable(y)
        y_ = net(x)
        loss = F.nll_loss(y_, y)
        optimizer.zero_grad()
        loss.backward()
        # update
        optimizer.step()
        if step % 200 == 0:
            print('Train Epoch: ', epoch, ' [', step * len(x), '/', len(train_loader.dataset), ' (', 100. * step / len(train_loader), '%)]\tLoss: ', loss.item())

def do_test():
    net.eval()
    test_loss = 0
    correct = 0
    # 测试集
    with torch.no_grad():
        for x, y in test_loader:
            x, y = Variable(x, volatile=True), Variable(y)
            y_ = net(x)
            # sum up batch loss
            test_loss += F.nll_loss(y_, y).item()
            # get the index of the max
            pred = y_.data.max(1, keepdim=True)[1]
            correct += pred.eq(y.data.view_as(pred)).cpu().sum()
        test_loss /= len(test_loader.dataset)
        print('\nTest set: Average loss: ', test_loss, ', Accuracy: ', int(correct), '/', len(test_loader.dataset),' (', 100. * int(correct) / len(test_loader.dataset),'%)\n')


for epoch in range(1,TRAINING_STEPS):
    do_train(epoch)
    do_test()

运行输出：