Pytorch-Mnist手写数字识别

import torch
import numpy as np
# 导入 Pytorch内置的mnist数据
from torchvision.datasets import mnist
# 导入预处理模块
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# 导入nn及优化器
import torch.nn.functional as F
import torch.optim as optim
from torch import nn

# 定义一些超参数
train_batch_size = 64
test_batch_size = 128
learning_rate = 0.01
num_epoches = 50
lr = 0.01
momentum = 0.5

# 下载数据并对数据进行预处理
# 定义预处理函数，这些预处理一次放在Compose函数中。
"""
root: 数据集，存在根目录processed/train.pt和processed/test.pt中
train: True = 训练集， False = 测试集
download: 如果为True，请从Internet下载数据集并将其放在根目录中，如果数据集已经下载，则不会再次下载
transform: 接收PIL影响并返回转换版本的函数/变换
target_transform: 一个接收目标并转换她的函数/变换
"""
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5],[0.5])])
# 下载数据集，并对数据集进行预处理
train_dataset = mnist.MNIST("./data", train=True, transform=transform, download=True)
test_dataset = mnist.MNIST("./data", train=True, transform=transform)
# dataloader是一个可迭代对象，可以使用迭代器一样使用
"""
dataset: 加载数据的数据集
batch_size: 加载批训练的训练个数
shuffle: 为True则在每个Epoch重新排列数据
sampler: 从数据集中提取样本
batch_sampler: 一次返回一批索引
num_worker: 用于数据加载的子进程数，0表示数据将在主进程中加载
collate_fn：合并样本列表以形成小批量
pin_memory: 如果为True，数据加载器在返回前张量复制到CUDA固定内存中。
drop_last: 如果数据集大小不能被batch_size整除，设置True可删除最后一个不完整的批处理。
如果设为False并且数据集的大小不能被batch_size整除, 则最后一个batch将更小。
"""
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
test_loader =DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)

import matplotlib.pyplot as plt

examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)

fig = plt.figure()
for i in range(6):
    plt.subplot(2, 3, i+1)
    plt.tight_layout()
    plt.imshow(example_data[i][0], cmap="gray", interpolation="none")
    plt.title("Ground Truth: {}".format(example_targets[i]))
    plt.xticks([])
    plt.yticks([])
plt.show()

# 构建网络
class Net(nn.Module):
    """
    使用Sequental构建网络，Sequential()函数功能是将网络的层组合到一起
    前馈神经网络由一个输入层、一个或多个隐藏层和一个输出层组成。
    """
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Net, self).__init__()
        self.layer1 = nn.Sequential(nn.Linear(in_dim, n_hidden_1), nn.BatchNorm1d(n_hidden_1))
        self.layer2 = nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2), nn.BatchNorm1d(n_hidden_2))
        self.layer3 = nn.Sequential(nn.Linear(n_hidden_2, out_dim), nn.BatchNorm1d(out_dim))
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = F.relu(self.layer3(x))
        return x
# 实例化网络
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 实例化网络
model = Net(28*28, 300, 100, 10)
model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

# 开始训练
losses = []
acces = []
eval_losses = []
eval_acces = []

for epoch in range(num_epoches):
    train_loss =0
    train_acc = 0
    model.train()
    # 动态修改参数学习率
    if epoch%5 ==0:
        optimizer.param_groups[0]['lr']*=0.1
    for img, label in train_loader:
        img = img.to(device)
        label = label.to(device)
        img = img.view(img.size(0),-1)
        c = img.size()
        # 前向传播
        out = model(img)
        loss = criterion(out, label)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 记录误差
        train_loss += loss.item()
        # 计算分类的准确率
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        train_acc += acc
    losses.append(train_loss / len(train_loader))
    acces.append(train_acc / len(train_loader))

    # 在测试集上检验效果
    eval_loss = 0
    eval_acc = 0
    # 将模型改为预测模型
    model.eval()
    for img, label in test_loader:
        img = img.to(device)
        label = label.to(device)
        img = img.view(img.size(0), -1)
        out = model(img)
        loss = criterion(out, label)
        # 记录误差
        eval_loss += loss.item()
        # 记录准确率
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        eval_acc += acc
    eval_losses.append(eval_loss / len(test_loader))
    eval_acces.append(eval_acc / len(test_loader))
    print("epoch:{}, Train loss: {:.4f}, Train Acc:{:.4f}, Test Loss:{:.4f}, Test acc:{}".format(epoch, train_loss / len(train_loader), train_acc / len(train_loader),
                                                                                                eval_loss / len(test_loader), eval_acc / len(test_loader)
                                                                                                ))