本文将利用之前学到的 PyTorch 的相关知识，建立一个全连接神经网络模型，用于识别手写字符。通过这个小demo完成数据集的预处理、数据加载器的生成、优化器的定义、损失的定义、全连接神经网络的搭建、训练与测试等过程。

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device

在下面代码中，会对所有数据变量添加一个 .to(device) 操作。如果当前环境支持 GPU 运行， .to(device) 就可以使变量转成可放入 GPU 中的类型。若不支持，.to(device) 就可以使变量转成可放入 CPU 中的类型。

一、import功能包

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

二、处理train、test数据集

加载 PyTorch 中的自带数据集合，该数据集合存在于torchvision.datasets 中，可以直接利用torchvision.datasets.MNIST获得。

batch_size=128
# 将数据集合下载到指定目录下,这里的transform表示，数据加载时所需要做的预处理操作
# 加载训练集合
train_dataset=torchvision.datasets.MNIST(
    root='.data',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=True)
# 加载测试集合
test_dataset=torchvision.datasets.MNIST(
    root='.data',
    train=False,
    transform=torchvision.transforms.ToTensor(),
    )

train_loader=torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True
        )

test_loader=torch.utils.data.DataLoader(
        dataset=test_dataset,
        batch_size=batch_size,
        shuffle=True
        )

3.查看几张图片

在定义完 PyTorch 能够识别的数据加载器后，我们可以加载几张图片，观察一下图片效果

import matplotlib.pyplot as plt
%matplotlib inline
# 加载测试集中的前 6 张图片
examples = iter(test_loader)
example_data, example_targets = examples.next()

for i in range(6):
    plt.subplot(2, 3, i+1)
    plt.imshow(example_data[i][0], cmap='gray')
plt.show()

输出如下：

四、model模型建立

模型的输入节点数和输出节点数为：

# 输入节点数就为图片的大小：28×28×1
input_size = 784
# 由于数字为 0-9，因此是10分类问题，因此输出节点数为 10
num_classes = 10
input_size, num_classes

第一层为输入层，节点数与图像大小相同。
第二次为隐藏层，节点数为 500 。
第三层为输出层，节点大小为 10 ，节点数大小和类别相同。

# 包含了一个隐含层的全联机神经网络
class NeuralNet(nn.Module):
    # 输入数据的维度，中间层的节点数，输出数据的维度
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.input_size = input_size
        self.l1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.l2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.l1(x)
        out = self.relu(out)
        out = self.l2(out)
        return out
# 建立了一个中间层为 500 的三层神经网络，且将模型转为当前环境支持的类型（CPU 或 GPU）
model = NeuralNet(input_size, 500, num_classes).to(device)
model

五、配置优化器

交叉熵损失函数：

criterion = nn.CrossEntropyLoss()

利用 PyTorch 中定义的梯度下降算法的Adam优化算法来进行模型的训练：

# 此时学习率为 0.001 ，你也可以根据实际情况，自行设置
learning_rate = 0.001
# 定义 Adam 优化器用于梯度下降
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

6.开始训练

模型训练的步骤可以说是固定的：
1.通过模型的正向传播，输出预测结果。
2.通过预测结果和真实标签计算损失。
3.通过后向传播，获得梯度。
4.通过梯度更新模型的权重。
5.进行梯度的清空。
6.循环上面操作，直到损失较小为止。

num_epochs = 5
# 数据总长度
n_total_steps = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 因为全连接会把一行数据当做一条数据，因此我们需要将一张图片转换到一行上
        # 原始数据集的大小: [100, 1, 28, 28]
        # 将每一张图片都转为一行向量，
        # resize 后的向量大小: [100, 784]
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)

        # 正向传播以及损失的求取
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播
        # 下面三句话固定：梯度清空，反向传播，权重更新
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print(
                f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')
print("模型训练完成")

放几张图，看看预测结果是否和真实一致：

# 测试样例
examples = iter(test_loader)
example_data, example_targets = examples.next()

# 图片的展示
for i in range(3):
    plt.subplot(1, 3, i+1)
    plt.imshow(example_data[i][0], cmap='gray')
plt.show()

# 结果的预测
images = example_data
images = images.reshape(-1, 28*28).to(device)
labels = labels.to(device)

# 正向传播以及损失的求取
outputs = model(images)
# 将 Tensor 类型的变量 example_targets 转为 numpy 类型的，方便展示
print("上面三张图片的真实结果：", example_targets[0:3].detach().numpy())
# 将得到预测结果
# 由于预测结果是 N×10 的矩阵，因此利用 np.argmax 函数取每行最大的那个值，最为预测值
print("上面三张图片的预测结果：", np.argmax(outputs[0:3].detach().numpy(), axis=1))

输出如下：

上面三张图片的真实结果： [4 1 5]
上面三张图片的预测结果： [4 1 5]

7.检查准确性

# I测试数据，计算模型的识别准确率
with torch.no_grad():
    n_correct = 0
    n_samples = 0
    for images, labels in test_loader:
        # 和训练代码一致
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)
        outputs = model(images)

        # 进行模型训练
        _, predicted = torch.max(outputs.data, 1)
        n_samples += labels.size(0)
        n_correct += (predicted == labels).sum().item()

    acc = 100.0 * n_correct / n_samples
    print(f'Accuracy of the network on the 10000 test images: {acc} %')

输出如下：

Accuracy of the network on the 10000 test images: 97.7 %