What is torch.nn really?

  • PyTorch通过精心设计的模块和类——torch.nn、torch.optim、Dataset及DataLoader——来协助构建和训练神经网络

  • 若要充分发挥其能力并针对具体问题实现定制化，就需要真正理解它们内部的运作机制

  • 为建立这种理解，我们将首先在不使用这些模型中任何功能的情况下，在MNIST数据集上训练一个基础神经网络；初始阶段仅使用最基本的PyTorch张量功能

  • 随后，逐步每次添加一个来自torch.nn、torch.optim、Dataset或DataLoader的功能组件，清晰展示每个部分的作用，以及它们如何使代码更简洁或更灵活

MNIST data setup

  • 使用经典的MNIST数据集，它由手绘数字（0到9之间）的黑白图像组成

  • 使用pathlib处理路径（属于Python 3标准库），并通过requests库下载数据集

1
2
3
4
5
6
7
8
9
10
11
12
13
14
from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"

PATH.mkdir(parents=True, exist_ok=True)

URL = "https://github.com/pytorch/tutorials/raw/main/_static/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():
        content = requests.get(URL + FILENAME).content
        (PATH / FILENAME).open("wb").write(content)

  • 数据集采用numpy array格式，并使用Python特有的数据序列化格式pickle进行存储

  • 每张图像的尺寸为28 x 28像素，并以长度为784（28x28）的扁平化行形式存储。需要先将其重塑为二维形式才能查看具体图像

1
2
3
4
5
6
7
8
9
10
11
12
from matplotlib import pyplot
import numpy as np

pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")
# ``pyplot.show()`` only if not on Colab
try:
    import google.colab
except ImportError:
    pyplot.show()
print(x_train.shape)

# out:(50000, 784)
  • PyTorch使用torch.tensor而非numpy arrays，因此我们需要对数据进行转换
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch

x_train, y_train, x_valid, y_valid = map(
    torch.tensor, (x_train, y_train, x_valid, y_valid)
)
n, c = x_train.shape
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

"""
out:
tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]]) tensor([5, 0, 4,  ..., 8, 4, 8])
torch.Size([50000, 784])
tensor(0) tensor(9)
"""

Neural net from scratch (without torch.nn)

  • 首先仅使用PyTorch张量操作来创建模型

  • PyTorch提供了生成随机或全零张量的方法，我们将利用这些方法为简单线性模型创建权重和偏置

  • 这些只是常规张量，但有一个特别重要的附加属性：告知PyTorch这些张量需要计算梯度。这将使PyTorch记录所有在张量上执行的操作，从而能够自动计算反向传播过程中的梯度

  • 对于权重参数，我们在初始化后设置requires_grad属性，因为不希望该初始化步骤被包含在梯度计算中（注意：PyTorch中尾部下划线表示原地操作）

  • 注：此处我们采用Xavier初始化方法（通过乘以1/sqrt(n)）来初始化权重

1
2
3
4
5
6

import math

weights = torch.randn(784, 10) / math.sqrt(784)
weights.requires_grad_()
bias = torch.zeros(10, requires_grad=True)

  • 得益于PyTorch自动计算梯度的能力，我们可以使用任何标准Python函数（或可调用对象）作为模型。只需编写基础的矩阵乘法和广播加法即可创建简单的线性模型

  • 同时还需要一个激活函数，这里我们将实现log_softmax并使用它

  • 注意：尽管PyTorch提供了大量预写的损失函数、激活函数等组件，但完全可以使用原生Python轻松编写自定义函数。PyTorch甚至会为函数自动生成快速加速器或向量化CPU代码

1
2
3
4
5
def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)

  • 在上述代码中，@符号表示矩阵乘法运算

  • 对单批数据（本例中为64张图像）调用该函数——这便是一次前向传播

  • 需要注意的是，由于初始权重为随机值，当前阶段的预测效果不会优于随机猜测

1
2
3
4
5
6
7
8
9
10
11
12
bs = 64  # batch size

xb = x_train[0:bs]  # a mini-batch from x
preds = model(xb)  # predictions
preds[0], preds.shape
print(preds[0], preds.shape)

"""
out:
tensor([-2.6859, -2.6669, -2.2688, -2.6470, -2.6484, -2.0377, -2.0170, -2.3224,
        -1.7795, -2.4400], grad_fn=<SelectBackward0>) torch.Size([64, 10])
"""

  • preds 张量不仅包含张量值，还包含一个梯度函数。稍后我们将使用它进行反向传播

  • 现在实现负对数似然作为损失函数（可以直接使用标准 Python）

1
2
3
4
def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll
  • 用随机模型检查损失值，以便观察反向传播后是否有所改善
1
2
3
4
5
6
7
yb = y_train[0:bs]
print(loss_func(preds, yb))

"""
out:
tensor(2.3058, grad_fn=<NegBackward0>)
"""
  • 还需要实现一个函数来计算模型的准确率。对于每个预测，如果最大值对应的索引与目标值相符，则预测正确
1
2
3
4
5
6
7
8
9
10
def accuracy(out, yb):
    preds = torch.argmax(out, dim=1)
    return (preds == yb).float().mean()

print(accuracy(preds, yb))

"""
out:
tensor(0.1094)
"""

  • 运行训练循环。每次迭代中，我们将：

    1. 选择一个minibatch的数据（大小为 bs）
    2. 使用模型进行预测
    3. 计算损失值
    4. loss.backward() 会更新模型的梯度（weights & bias）

  • 现在我们使用这些梯度来更新权重和偏置。我们在 torch.no_grad() 上下文管理器中执行此操作，因为我们不希望这些操作被记录到下一次梯度计算中

  • 随后将梯度归零，为下一个循环做好准备。否则梯度会持续累加所有已发生的操作记录（即 loss.backward() 会在已有存储值的基础上累加梯度，而不是替换它们）

  • 提示：可以使用标准 Python 调试器逐步执行 PyTorch 代码，从而检查每个步骤中的各种变量值

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        #         set_trace()
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()

  • 至此，我们已经从零开始创建并训练了一个极简神经网络（本例中为 logistic regression，因为没有隐藏层）

  • 现在检查 loss 和 accuracy 并与之前的结果进行对比。预期 loss 应该下降而 accuracy 应该上升，实际结果确实如此

1
2
3
4
5
6
print(loss_func(model(xb), yb), accuracy(model(xb), yb))

"""
out:
tensor(0.0835, grad_fn=<NegBackward0>) tensor(1.)
"""

Using torch.nn.functional

  • 现在将对代码进行重构，使其保持原有功能不变，但开始利用PyTorch的nn类来使代码更简洁灵活

  • 第一步通过用torch.nn.functional（通常按惯例导入为F）中的函数替代我们手写的激活函数和损失函数来缩短代码。该模块包含torch.nn库中的所有函数

  • 除了各种损失函数和激活函数外，还可以在这里找到一些用于创建神经网络的便捷函数，例如池化函数（此外还有用于卷积、线性层等操作的函数）

  • 如果使用的是负对数似然损失和log softmax激活函数，那么PyTorch提供了一个结合了这两者的单一函数F.cross_entropy。这样甚至可以从模型中移除激活函数

1
2
3
4
5
6
7
8
9
10
11
12
13
import torch.nn.functional as F

loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

"""
out:
tensor(0.0838, grad_fn=<NllLossBackward0>) tensor(1.)
"""

Refactor using nn.Module

  • 接下来使用nn.Module和nn.Parameter来构建更清晰简洁的训练循环。通过子类化nn.Module（它本身是一个能够跟踪状态的类）来实现

  • 在这种情况下需要创建一个包含权重、偏置和前向传播方法的类

  • nn.Module具有许多我们将用到的属性和方法（例如.parameters()和.zero_grad()）

  • 注意：nn.Module（大写M）是PyTorch特有的概念，是一个我们会频繁使用的类。不要将nn.Module与Python中（小写m）module的概念混淆，后者是指可被导入的Python代码文件

1
2
3
4
5
6
7
8
9
10
from torch import nn

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(784, 10) / math.sqrt(784))
        self.bias = nn.Parameter(torch.zeros(10))

    def forward(self, xb):
        return xb @ self.weights + self.bias
  • 由于我们现在使用对象而不是单纯使用函数，因此首先需要实例化我们的模型：
1
model = Mnist_Logistic()

  • 现在我们可以用与之前相同的方式计算损失

  • 注意：nn.Module对象的使用方式与函数类似（即可调用），但底层PyTorch会自动调用我们的forward方法

1
2
3
4
5
6
print(loss_func(model(xb), yb))

"""
out:
tensor(2.4364, grad_fn=<NllLossBackward0>)
"""
  • 在先前的训练循环中，我们必须按名称逐个更新每个参数的值，并手动将每个参数的梯度单独清零
1
2
3
4
5
with torch.no_grad():
    weights -= weights.grad * lr
    bias -= bias.grad * lr
    weights.grad.zero_()
    bias.grad.zero_()
  • 现在我们可以利用 model.parameters() 和 model.zero_grad()（这两个方法都是由 PyTorch 为 nn.Module 定义的）来简化这些步骤，并降低遗漏某些参数的错误风险，尤其是在模型更为复杂的情况下：
1
2
3
with torch.no_grad():
    for p in model.parameters(): p -= p.grad * lr
    model.zero_grad()
  • 将把训练循环封装到fit函数中，以便后续可以重复运行
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def fit():
    for epoch in range(epochs):
        for i in range((n - 1) // bs + 1):
            start_i = i * bs
            end_i = start_i + bs
            xb = x_train[start_i:end_i]
            yb = y_train[start_i:end_i]
            pred = model(xb)
            loss = loss_func(pred, yb)

            loss.backward()
            with torch.no_grad():
                for p in model.parameters():
                    p -= p.grad * lr
                model.zero_grad()

fit()

print(loss_func(model(xb), yb))

"""
out:
tensor(0.0815, grad_fn=<NllLossBackward0>)
"""

Refactor using nn.Linear