Pytorch自定义网络结构训练手写数字识别模型

数据集

MINIST数据集，可以自己在网上下载

训练工具

华为云的modelarts中的notebook，华为云近期有免费的P100服务器，供开发者使用，虽然每次只能用一个小时，但训练这种简单的模型已经够用了。

训练代码

导入使用的工具包

import torch as t
import torchvision
from torchvision import datasets , transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import pylab

use_gpu = t.cuda.is_available()
print(use_gpu)

#Compose是一个容器，传入的参数是列表，ToTensor()，类型变换，Normalize是数据标准化，去均值，除标准差
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize(mean=[0.5],
                                                     std=[0.5])])
# Normalize 数据标准化变换  mean和std为原始数据的均值和标准差（此处直接赋值）

# 首先获取手写数字的训练集和测试集
# root 用于指定数据集在下载之后的存放路径
# transform 用于指定导入数据集需要对数据进行那种变化操作
# train是指定在数据集下载完成后需要载入那部分数据，
# 如果设置为True 则说明载入的是该数据集的训练集部分
# 如果设置为FALSE 则说明载入的是该数据集的测试集部分
data_train = datasets.MNIST(root = "../data/",
                            transform = transform,
                            train =True,
                            download = True)
data_test = datasets.MNIST(root = "../data/",
                            transform = transform,
                            train =False)

#数据预览和数据装载
# 下面对数据进行装载，我们可以将数据的载入理解为对图片的处理，
# 在处理完成后，我们就需要将这些图片打包好送给我们的模型进行训练 了  而装载就是这个打包的过程
# dataset 参数用于指定我们载入的数据集名称
# batch_size参数设置了每个包中的图片数据个数
#  在装载的过程会将数据随机打乱顺序并进打包
data_loader_train = t.utils.data.DataLoader(dataset = data_train,
                                            batch_size = 64,
                                            shuffle = True
                                            )

data_loader_test = t.utils.data.DataLoader(dataset = data_test,
                                            batch_size = 64,
                                            shuffle = True  #将顺序随机打乱
                                            )


class Model(t.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.convl = t.nn.Sequential(
            t.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),  # 卷积层
            # 1*28*28
            # 输入通道数，输出通道数 都为整形
            # kernel_size指卷积核的大小；stride指步长，即卷积核或者pooling窗口的滑动位移。
            # padding指对input的图像边界补充一定数量的像素，目的是为了计算位于图像边界的像素点的卷积响应；
            # ( input_size + 2*padding - kernel_size ) / stride+1 = output_size
            t.nn.ReLU(),
            t.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),  # 卷积层

            t.nn.ReLU(),
            t.nn.MaxPool2d(stride=2, kernel_size=2)
        )
        self.dense = t.nn.Sequential(  # 全连接层
            t.nn.Linear(14*14*128, 1024),
            t.nn.ReLU(),
            t.nn.Dropout(p=0.5),  # 防止训练过程发生过拟合，torch.nn.Dropout对所有元素中每个元素按照概率0.5更改为零
            t.nn.Linear(1024, 10)
        )

    def forward(self, x):  # 前向传播
        x = self.convl(x)  # 输入卷积层
        x = x.view(-1, 14 * 14 * 128)  # torch里面，view函数相当于numpy的reshape
        # -1表示一个不确定的数，就是你如果不确定你想要reshape成几行，但是列确定，-1会自动更改为合适的值
        x = self.dense(x)  # 输入全连接层
        return x

model = Model()
cost = t.nn.CrossEntropyLoss()
optimizer = t.optim.Adam(model.parameters())
if use_gpu:
    model = model.cuda()
    cost = cost.cuda()
#print(model)

n_epochs = 5
for epoch in range(n_epochs):
    running_loss = 0.0
    running_correct = 0.0
    print("Epoch {}/{}".format(epoch,n_epochs))
    print("-"*10)
    for data in data_loader_train:
        X_train, y_train = data
        if (use_gpu):
            X_train, y_train = X_train.cuda(), y_train.cuda()
        X_train, y_train = Variable(X_train), Variable(y_train)
        output = model(X_train)
        _, pred = t.max(output.data, 1)
        optimizer.zero_grad()
        loss = cost(output, y_train)
        #print("[l]:{:.4f}".format(loss))
        loss.backward()
        optimizer.step()
        running_loss += loss.data
        running_correct += t.sum(pred == y_train.data)
    test_correct = 0
    for data in data_loader_test:
        X_test, y_test = data
        if (use_gpu):
            X_test, y_test = X_test.cuda(), y_test.cuda()
        X_test, y_test = Variable(X_test), Variable(y_test)
        output = model(X_test)
        _, pred = t.max(output.data, 1)
        test_correct += t.sum(pred == y_test.data)
    print("Loss is:{:.4f}, Train_accuracy is {:.4f}%, Test_accuracy is {:.4f}%"
          .format(running_loss/len(data_train),100*running_correct/len(data_train), 100*test_correct/len(data_test)))

t.save(model, 'model.pkl')
#model = torch.load('model.pkl')