手写数字识别源代码+注释

import torch
import torchvision
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
import math
import torch
import torch.nn as nn

transform = transforms.Compose([
    transforms.ToTensor(),  # 将PILImage转换为张量
    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),  # 将下载好的图片从单通道repeat成3通道RGB图片
])
# 下载的图片集的格式

data_train = datasets.MNIST(root="./data/",
                            transform=transform,
                            train=True,
                            download=True)

data_test = datasets.MNIST(root="./data/",
                           transform=transform,
                           train=False)
# 下载图片集

data_loader_train = torch.utils.data.DataLoader(dataset=data_train,
                                                batch_size=64,
                                                shuffle=True,
                                                )
data_loader_test = torch.utils.data.DataLoader(dataset=data_test,
                                               batch_size=64,
                                               shuffle=True,
                                               )
# 将下载好的图片集进行打乱分类制成新的数据集

# images, labels = next(iter(data_loader_train))
# print(images.shape)
# img = torchvision.utils.make_grid(images)  # 经过torchvision.utils.make_grid后图片维度变为channel,h,w三维
# img = img.numpy().transpose(1, 2, 0)  # 转化为np格式并且调换w和channel维度
# print([labels[i] for i in range(64)])
# plt.imshow(img)
# 用网格函数对图片进行排列并展示图片

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(stride=2, kernel_size=2)
        )
        self.dense = nn.Sequential(
            nn.Linear(14*14*128, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),  # 随机归零，放置过拟合
            nn.Linear(1024, 10)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x


model = Model()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

epoch_n = 5
for epoch in range(epoch_n):

    # 训练开始
    running_loss = 0.0
    running_correct = 0.0
    print("Epoch {}/{}".format(epoch, epoch_n))
    for data in data_loader_train:

        X_train, Y_train = data
        X_train, Y_train = Variable(X_train), Variable(Y_train)  # 变成张量

        outputs = model(X_train)  # 训练一次

        _, pre = torch.max(outputs.data, 1)   # max函数返回最大值和下标，取下标

        optimizer.zero_grad()  # 把梯度置零

        loss = loss_fn(outputs, Y_train)  # 根据损失函数得到代价
        loss.backward()  # 代价反向传播

        optimizer.step()  # 更新参数

        running_loss += loss.data  # 更新损失值
        running_correct += torch.sum(pre == Y_train.data)  # 更新正确率

    # 测试开始
    testing_correct = 0.0
    for data in data_loader_test:

        X_train, Y_train = data
        X_train, Y_train = Variable(X_train), Variable(Y_train)  # 变成张量

        outputs = model(X_train)  # 训练一次

        _, pre = torch.max(outputs.data, 1)  # max函数返回最大值和下标，取下标

        testing_correct += torch.sum(pre == Y_train.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*testing_correct/len(data_test)))