算法工程师的修养 | PyTorch 的 nn 详解-技术圈

来源：知乎—zideajang

地址：https://zhuanlan.zhihu.com/p/387133209、https://zhuanlan.zhihu.com/p/387193068

PyTorch 提供了设计优雅的模块和类 torch.nn、torch.optim、Dataset 和 DataLoader 来帮助你创建和训练神经网络。可以根据您的具体问题对这些模块进行定制，充分发挥这些模块的威力。那么就需要真正理解这些模块到底在什么。为了帮助您更好理解这些模块，我们将首先在 MNIST 数据集上训练基本的神经网络，而不使用这些模型的任何特征。最初将只使用最基本的 PyTorch tensor 功能。然后，我们将逐步从torch.nn、torch.optim、Dataset 或 DataLoader 中添加一个特征，准确地展示每一模块的作用，以及如何使代码更简洁或更灵活。

对您的一点要求

需要对 PyTorch 这个框架有所了解，了解 python 语法可以 python 写点东西，并且熟悉 tensor 操作的基础知识。《》,如果熟悉 Numpy 数组操作，即使对 tensor 并不了解也没关系，因为你会发现这里使用的 PyTorchd 对 tensor 操作几乎是一样的。

用 Tensor 白手起家实现一个网络

准备数据集

这里还是用经典的 MNIST 数据集，如果你对这数据集还不熟悉，说明你应该是刚刚接触深度学习没几天。MNIST 数据集是由手写数字（0到9之间）的黑白图像所组成。

首先需要获取数据集，将使用 pathlib 来处理下载路径(Python 3标准库的一部分)，使用 request 来下载数据集。

python from pathlib import Path import requestsDATA_PATH = Path("data") PATH = DATA_PATH / "mnist"PATH.mkdir(parents=True, exist_ok=True)URL = "https://github.com/pytorch/tutorials/raw/master/_static/" FILENAME = "mnist.pkl.gz"if not (PATH / FILENAME).exists():         content = requests.get(URL + FILENAME).content         (PATH / FILENAME).open("wb").write(content)

这个数据集是一个 numpy 数组，并使用 pickle 方式进行存储，这是一种python特有的数据序列化格式。

pythonimport pickleimport gzip
with gzip.open((PATH / FILENAME).as_posix(), "rb") as f:        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")```
每张图片尺寸是 28 x 28，被存储为长度为 784（=28x28）的向量。可以利用 `matplotlib` 看看其中一张图像，想要将其展示出来首先还需要先把其 reshape 为2D。
```pythonfrom matplotlib import pyplotimport numpy as np
pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")print(x_train.shape)

(50000, 784)

将 numpy 数组转换为 Tensor


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.shapeprint(x_train, y_train)print(x_train.shape)print(y_train.min(), y_train.max())

输出


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)
开始搭建网络(不借助外力 nn 什么的)

让我们首先创建一个模型，只使用PyTorch的张量操作。我们假设你已经熟悉了神经网络的基础知识。如果你不熟悉，你可以在course.fast.ai上学习）。

PyTorch 提供多种多样创建 tensor 的方式，创建随机数的 tensor 或者创建一个全 0 的 tensor《》，现在我们就为一个简单的线性模型创建的权重和偏置。创建参数只用到了普通的 tensor，创建 tensor 时我们可以显式告诉 PyTorch 这些 tensor 是否需要计算梯度。如果需要计算梯度，则 PyTorch 将对 tensor 在前向传播所做的所有操作进行存储，这样就可以在反向传播过程中对其自动地进行梯度计算了。

如果是权重，在初始化后设置 requirest_grad，requires_grad() 将更改 tensor 是否求导的状态并返回。请注意，PyTorch 中的函数名以结尾的，表示操作是在 pytorch 内部进行的)。

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 实现一些简单损失函数和激活函数。《激活函数》

def log_softmax(x):    return x - x.exp().sum(-1).log().unsqueeze(-1)
def model(xb):    return log_softmax(xb @ weights + bias)

在上面，@ 表示矩阵内积操作。我们将在一批数据上调用我们的函数（在本例中，64张图片）。这就是一个前向传递。请注意，我们的预测在这个阶段不会比随机好，因为我们从随机权重开始。

bs = 64  # batch size
xb = x_train[0:bs]  # a mini-batch from xpreds = model(xb)  # predictionspreds[0], preds.shapeprint(preds[0], preds.shape)tensor([-1.9398, -2.3529, -2.2999, -2.6261, -2.6767, -1.3650, -2.3081, -2.7904,        -2.9199, -3.0043], grad_fn=<SelectBackward>) torch.Size([64, 10])

正如你所看到的，preds 张量不仅包含数值，而且还包含一个梯度函数，这个 SelectBackward 表示该 tensor 是如何获得 SelectBackward 是计算图的边，tensor 通常是计算图的节点。

让我们实现负对数似然来作为损失函数（可以直接使用标准 Python 来实现）。《》

def nll(input, target):    return -input[range(target.shape[0]), target].mean()
loss_func = nll

用随机数作为参数初始化一个模型来验证一下损失函数是否好用，损失函数随后反向传播更新参数的一个方向，都是朝着让损失函数变小。

yb = y_train[0:bs]
print(loss_func(preds, yb))
tensor(2.3554, grad_fn=<NegBackward>)

让我们也实现一个函数来计算我们模型的准确性。对于每个预测，如果数值最大的索引与目标值相匹配，那么预测就是正确的。


def accuracy(out, yb):    preds = torch.argmax(out, dim=1)    return (preds == yb).float().mean()print(accuracy(preds, yb))tensor(0.1562)

我们现在可以进行训练。在训练迭代过程中，对于每次迭代，我们将。

选择一个小批量的数据（大小为bs）。
使用模型进行预测
计算损失
loss.backward()更新模型参数的梯度，在这里是指权重和偏置。

现在使用这些梯度来更新权重和偏置。我们在 torch.no_grad() 下进行用梯度来更新参数，因为不希望这些行为被记录在我们下一次计算梯度。然后将梯度设置为 0，这样就可以为下一个迭代做好准备。否则，梯度将记录所有已经发生的操作。

from IPython.core.debugger import set_trace
lr = 0.5  # learning rateepochs = 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_()

就这样：我们完全从头开始创建并训练了一个最小的神经网络（在这种情况下，是一个逻辑回归，因为我们没有隐藏层）！让我们检查一下损失和准确率，并与我们之前得到的数据进行比较。

现在来通过损失值和准确率来验证以下从数据上来看，这些指标与之前得到的数据进行比较，感觉损失值有了显著减小，准确率也得到提高。

print(loss_func(model(xb), yb), accuracy(model(xb), yb))tensor(0.0815, grad_fn=<NegBackward>) tensor(1.)

利用 nn 重构神经网络

使用 torch.nn.functional 重构损失函数和激活函数

现在，利用 nn 模块提供功能来重构上面的代码，使用 PyTorch 的 nn 类的好处是，使代码更加简洁和灵活。

也是最简单的一步，就是用 torch.nn.functional(习惯上我们都会给这个 torch.nn.functional 一个别名 F )，中提供函数替换之前自己实现的的激活和损失函数。这个模块包含了 torch.nn 库中的所有函数。除了提供常用大部分的损失和激活函数外，还会在这里找到一些用于创建神经网络的便捷函数，可以使用这些功能来实现如池化函数，(甚至还可以实现卷积函数、线性层等，但是通常我们不会这么做，因为在 pytorch 其他模块已经这些复杂复用性强的层进行一步进行封装供我们调用)。

对于上面自己实现负对数似然损失函数和 softmax 激活函数，其实在 Pytorch 中已经提供了一个函数 F.cross_entropy 将两者来结合起来供开发者使用，所以。

import torch.nn.functional as F
loss_func = F.cross_entropy
def model(xb):    return xb @ weights + biasprint(loss_func(model(xb), yb), accuracy(model(xb), yb))tensor(0.0815, grad_fn=<NllLossBackward>) tensor(1.)

使用 nn.Module 重构神经网络

接下来，我们将使用 nn.Module 和 nn.Parameter 让训练循环的代码得更清晰、更简洁的。nn.Module 就是一个类，能够保持对状态的跟踪，我们可以创建 nn.Module 子类，nn.Module有很多属性和方法（如.parameters()和.zero_grad()），继承后将使用这些属性和方法。

nn.Module（大写 M）是 PyTorch 特有的概念，也是将经常使用的一个类。nn.Module 不能与 Python 中的（小写m）module 概念混淆，后者(module)是一个可以导入的 Python 代码文。

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.biasmodel = Mnist_Logistic()print(loss_func(model(xb), yb))tensor(2.3154, grad_fn=<NllLossBackward>)

以前，对于我们的训练循环，必须按名称更新每个参数的值，并手动为每个参数进行重置为 0，代码如下


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

将把的训练循环封装在一个拟合函数中


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))

tensor(0.0830, grad_fn=<NllLossBackward>)

利用 nn.Linear 进行重构

利用 nn.Linear 进行重构

继续重构我们的代码，这次不再手动定义和初始化 self.weights 和 self.bias，在前向传播函数定义函数xb @ self.weights + self.bias，而是使用 Pytorch 提供的 nn.Linear 线性层来实现同等功能。Pytorch 有很多类型的预定义层，这样可以大大简化我们的代码量，加快开发速度。

class Mnist_Logistic(nn.Module):    def __init__(self):        super().__init__()        self.lin = nn.Linear(784, 10)
    def forward(self, xb):        return self.lin(xb)

实例化定义好的模型，依旧用同样方法检测一下

model = Mnist_Logistic()print(loss_func(model(xb), yb))tensor(2.3734, grad_fn=<NllLossBackward>)fit()
print(loss_func(model(xb), yb))tensor(0.0822, grad_fn=<NllLossBackward>)

利用 optim 进行重构优化器

Pytorch也有一个包含各种优化算法的模块，torch.optim。使用优化器中的 step 方法在每次迭代更新参数，而不是手动更新每个参数。

from torch import optimdef get_model():    model = Mnist_Logistic()    return model, optim.SGD(model.parameters(), lr=lr)
model, opt = get_model()print(loss_func(model(xb), yb))
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()        opt.step()        opt.zero_grad()
print(loss_func(model(xb), yb))tensor(2.3110, grad_fn=<NllLossBackward>)tensor(0.0816, grad_fn=<NllLossBackward>)

使用 Dataset 进行重构

PyTorch 有一个抽象的数据集类，

len 函数: （由Python的标准len函数调用）
getitem函数: 接受索引参数，然后根据索引返回一个索引对应的数据

PyTorch 的 TensorDataset 是一个包裹 tensor 的数据集。通过定义长度和索引方式获取数据，有点类似我们熟悉的迭代器


from torch.utils.data import TensorDatasetmodel, opt = get_model()
for epoch in range(epochs):    for i in range((n - 1) // bs + 1):        xb, yb = train_ds[i * bs: i * bs + bs]        pred = model(xb)        loss = loss_func(pred, yb)
        loss.backward()        opt.step()        opt.zero_grad()
print(loss_func(model(xb), yb))tensor(0.0819, grad_fn=<NllLossBackward>)

使用 DataLoader 重构数据加载器

Pytorch 的 DataLoader 将数据成批提供给模型。对于数据集可以创建一个 DataLoader 。DataLoader 使得好处迭代批处理数据变得更加容易。替换这部分代码 train_ds[i*bs : i*bs+bs], 自动成批输出数据工作交给 DataLoader 完成

from torch.utils.data import DataLoader
train_ds = TensorDataset(x_train, y_train)train_dl = DataLoader(train_ds, batch_size=bs)from torch.utils.data import DataLoader
train_ds = TensorDataset(x_train, y_train)train_dl = DataLoader(train_ds, batch_size=bs)model, opt = get_model()
for epoch in range(epochs):    for xb, yb in train_dl:        pred = model(xb)        loss = loss_func(pred, yb)
        loss.backward()        opt.step()        opt.zero_grad()
print(loss_func(model(xb), yb))

输出

tensor(0.0822, grad_fn=<NllLossBackward>)

的确要感谢以下 Pytorch 为我们这么贴心提供了 nn.Module、nn.Parameter、Dataset 和 DataLoader 好用方便的模块，大大减少我们开发的工作量。

添加验证

上面我们通过使用 PyTorch 为我们提供的模块对重构了神经网络构建和训练这部分代码，接下来我们进一步完善代码和功能，因为缺少验证集，所以总感觉缺少点什么似的，我们现在就是在训练过程中添加验证这个步骤。

对训练数据顺序进行重排对于防止批次之间的关联性和过度拟合是很重要的。另一方面，无论是否对验证集进行洗牌，验证损失值都是相同的，所以就没有必要花费额外的时间对验证数据进行洗牌。

将对验证集使用的批处理的量是训练集的 2 倍。这是因为验证集不需要反向传播，因此占用的内存较少(因为不需要存储梯度)。利用这一点，可以将验证集批量设置大一些，这样可以更快速地计算损失值。

train_ds = TensorDataset(x_train, y_train)train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)
valid_ds = TensorDataset(x_valid, y_valid)valid_dl = DataLoader(valid_ds, batch_size=bs * 2)

我们将在每个 epoch 迭代完成时计算并输出验证损失值。

值的注意的是，总是在训练前调用 model.train()，在推理前使用 model.eval()，因为在训练和推理阶段网络中的 nn.BatchNorm2d 和 nn.Dropout 等层行为不同，以确保这些不同阶段的适当行为所以在不同阶段调用不同函数。


model, opt = get_model()
for epoch in range(epochs):    model.train()    for xb, yb in train_dl:        pred = model(xb)        loss = loss_func(pred, yb)
        loss.backward()        opt.step()        opt.zero_grad()
    model.eval()    with torch.no_grad():        valid_loss = sum(loss_func(model(xb), yb) for xb, yb in valid_dl)
    print(epoch, valid_loss / len(valid_dl))0 tensor(0.3134)1 tensor(0.4116)