如今，我们已经了解了 PyTorch 中张量及其运算，但这远远不够。本次实验将学会如何使用 PyTorch 方便地构建神经网络模型，以及 PyTorch 训练神经网络的步骤及方法。
知识点
- PyTorch 构建神经网络
- Sequential 容器结构
- 使用 GPU 加速训练
- 模型保存与推理
PyTorch 构建神经网络
方便定义不同类型的 Tensor 及利于反向传播的 Autograd 机制是深度学习框架的重要特点，但真正带来极大便利的，莫过于已经封装好的不同神经网络结构组件，包括不同类型的层以及各式各样的损失函数、激活函数、优化器等。
PyTorch 搭建神经网络结构的组件在 torch.nn 中 ，这些神经网络层大多以类出现，例如全连接层：torch.nn.Linear() ，MSE 损失函数类：torch.nn.MSELoss() 等。
除此之外，torch.nn.functional 下面同样也有神经网络层，激活函数，损失函数等，但均以函数出现，例如全连接层函数：torch.nn.functional.linear() ，MSE 损失函数：torch.nn.functionalmse_loss() 等。
总之，torch.nn 下面包含了神经网络组件类（大写字母），而 torch.nn.functional 包含了神经网络组件函数（小写字母）。
本次实验使用的数据集为 MNIST，你可以把它看成是 DIGITS 数据集的增强版，都是手写字符任务。前面我们使用过 Fashion MNIST 数据集，MNIST 数据集和其样本特征是一致的，只是样本类别不一样。MNIST 中每个样本是 28×28 的矩阵，目标是字符 0-9。
我们可以直接使用 PyTorch 提供的计算机视觉增强模块 torchvision 加载 MNIST 数据集。由于数据集托管在外网服务器上，国内的下载速度较慢，实验从国内服务器下载该数据集。
# 从国内服务器下载数据集
!wget -nc "http://labfile.oss.aliyuncs.com/courses/1081/MNIST.zip"
!unzip -o "MNIST.zip"
import torchvision
# 加载训练数据，参数 train=True，供 60000 条
train = torchvision.datasets.MNIST(
root='.', train=True, transform=torchvision.transforms.ToTensor(), download=True)
# 加载测试数据，参数 train=False，供 10000 条
test = torchvision.datasets.MNIST(
root='.', train=False, transform=torchvision.transforms.ToTensor(), download=True)
上面的代码中，transform=torchvision.transforms.ToTensor() 是利用了 torchvision 提供的 transforms 直接将原 NumPy 数组转换为 PyTorch 张量。现在，你可以尝试输出训练和测试数据的特征和目标查看。
train.data.shape, train.targets.shape, test.data.shape, test.targets.shape
接下来，我们还需要使用 PyTorch 提供的一个组件对数据进行封装。torch.utils.data.DataLoader 是 PyTorch 提供的及其常用的数据加载器，它可以将数据集封装成迭代器以方便我们后续进行小批量加载，数据打乱等操作。数据加载器准备好之后，后续只需要通过 for 循环来使用即可。
import torch
# 训练数据打乱，使用 64 小批量
train_loader = torch.utils.data.DataLoader(dataset=train,
batch_size=64, shuffle=True)
# 测试数据无需打乱，使用 64 小批量
test_loader = torch.utils.data.DataLoader(dataset=test,
batch_size=64, shuffle=False)
train_loader, test_loader
接下来，我们学习 PyTorch 构建神经网络的经典方法，也是官方推荐使用的一种方法。
首先，torch.nn 中的一个基础类 torch.nn.Module 。该类是 PyTorch 中所有神经网络的基类，它既可以表示神经网络中的某层，也可以表示若干层的神经网络。torch.nn 中的各个类实际上就是由 torch.nn.Modules 继承而拓展。所以，在实际使用中，我们可以继承nn.Module，撰写自定义网络层。
所以，当我们搭建神经网络时，也需要继承 torch.nn.Module。我们准备搭建一个包含两个隐含层的全连接网络。
输入（784） → 全连接层 1 （784, 512）→ 全连接层 2 （512, 128）→ 输出（10）
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(784, 512) # 784 是因为训练是我们会把 28*28 展平
self.fc2 = nn.Linear(512, 128) # 使用 nn 类初始化线性层（全连接层）
self.fc3 = nn.Linear(128, 10)
def forward(self, x):
x = F.relu(self.fc1(x)) # 直接使用 relu 函数，也可以自己初始化一个 nn 下面的 Relu 类使用
x = F.relu(self.fc2(x))
x = self.fc3(x) # 输出层一般不激活
return x
我们定义了新的神经网络结构类 Net()，并使用 nn.Linear 组合了 3 个线性层（全连接层）。前向传播过程中，代码使用了 PyTorch 中常用的函数模块 torch.nn.functional 提供的 RELU 激活函数，实际上你也可以通过实例化 nn.Relu 来达到同样的效果。
下面，我们实例化自定义神经网络类：
model = Net()
model
PyTorch 的好处在于你可以初始化一个 784 长度的随机值样本传入网络，测试一下输出：
model(torch.randn(1, 784))
目前，我们已经搭好了前向传播网络。下面的步骤和 TensorFlow 非常相似：定义损失函数，优化器以及开始训练。
loss_fn = nn.CrossEntropyLoss() # 交叉熵损失函数
opt = torch.optim.Adam(model.parameters(), lr=0.002) # Adam 优化器
这里，我们选择了十分常用的交叉熵损失函数 nn.CrossEntropyLoss ，以及 Adam 优化器 torch.optim.Adam 。值得注意的是，PyTorch 中优化器需传入模型的参数 model.parameters()，这是 PyTorch 的一个使用特性。
接下来，我们就可以开始训练了，这部分代码非常重要。
def fit(epochs, model, opt):
print("Start training, please be patient.")
# 全数据集迭代 epochs 次
for epoch in range(epochs):
# 从数据加载器中读取 Batch 数据开始训练
for i, (images, labels) in enumerate(train_loader):
images = images.reshape(-1, 28*28) # 对特征数据展平，变成 784
labels = labels # 真实标签
outputs = model(images) # 前向传播
loss = loss_fn(outputs, labels) # 传入模型输出和真实标签
opt.zero_grad() # 优化器梯度清零，否则会累计
loss.backward() # 从最后 loss 开始反向传播
opt.step() # 优化器迭代
# 自定义训练输出样式
if (i+1) % 100 == 0:
print('Epoch [{}/{}], Batch [{}/{}], Train loss: {:.3f}'
.format(epoch+1, epochs, i+1, len(train_loader), loss.item()))
# 每个 Epoch 执行一次测试
correct = 0
total = 0
for images, labels in test_loader:
images = images.reshape(-1, 28*28)
labels = labels
outputs = model(images)
# 得到输出最大值 _ 及其索引 predicted
_, predicted = torch.max(outputs.data, 1)
correct += (predicted == labels).sum().item() # 如果预测结果和真实值相等则计数 +1
total += labels.size(0) # 总测试样本数据计数
print('============ Test accuracy: {:.3f} ============='.format(
correct / total))
fit(epochs=1, model=model, opt=opt) # 训练 1 个 Epoch，预计持续 10 分钟
上方的代码中有详细的注释，但依旧有几点值得注意的地方。
首先，由于 PyTorch 没有提供类似于 Flatten 这样的展平类，所以我们通过 reshape 操作将输入 28×28 展平为 784，使其和网络结构参数符合。你也可以使用 view，但官方更推荐使用 reshape 。
其次，opt.zero_grad() 这步非常关键。由于 PyTorch 设计时梯度会累计，所以我们需要手动清零以实现传入一个 Batch，计算梯度，然后更新参数，从而不会因为前面的梯度累计影响后面的参数更新。但 PyTorch 这样设计也是有原因的，比如当我们想提升 Batch 的大小而硬件又无法处理较多数据时，就可以通过梯度累积机制，等待传入多个 Batch 后再更新参数并执行清零，这就给了开发更多的灵活性。同时，后续循环神经网络中也可能利用到这个特性。