torch.nn 到底是什么？

作者: Jeremy Howard, fast.ai。感谢 Rachel Thomas 和 Francisco Ingham。

我们建议将此教程作为笔记本而非脚本来运行。要下载笔记本 (.ipynb) 文件，请点击页面顶部的链接。

PyTorch 提供了精心设计的模块和类，如 torch.nn、torch.optim、Dataset 和 DataLoader，以帮助您创建和训练神经网络。为了充分利用它们的功能并根据您的问题进行定制，您需要真正理解它们的具体作用。为了培养这种理解，我们将首先在 MNIST 数据集上训练一个基本的神经网络，而不使用这些模型中的任何功能；我们最初只会使用最基本的 PyTorch 张量功能。然后，我们将逐步添加 torch.nn、torch.optim、Dataset 或 DataLoader 中的一个功能，逐一展示每个部分的作用，以及它如何使代码更简洁或更灵活。

本教程假设您已经安装了 PyTorch，并且熟悉张量操作的基础知识。（如果您熟悉 Numpy 数组操作，您会发现这里使用的 PyTorch 张量操作几乎相同）。

MNIST 数据准备

我们将使用经典的 MNIST 数据集，该数据集包含手绘数字（0 到 9 之间）的黑白图像。

我们将使用 pathlib 来处理路径（Python 3 标准库的一部分），并使用 requests 下载数据集。我们只会在使用时导入模块，因此您可以清楚地看到每个步骤中使用了哪些模块。

frompathlibimport Path
importrequests

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数组格式，并使用pickle（一种特定于Python的数据序列化格式）进行存储。

importpickle
importgzip

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) 的扁平化行存储。让我们来看一下其中一张图片；首先需要将其重塑为二维形式。

frommatplotlibimport pyplot
importnumpyasnp

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

nn tutorial

(50000, 784)

PyTorch 使用 torch.tensor 而不是 numpy 数组，因此我们需要转换数据。

importtorch

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

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)

从零开始构建神经网络（不使用 `torch.nn`）

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

PyTorch 提供了创建随机或零填充张量的方法，我们将使用这些方法来为简单的线性模型创建权重和偏置。这些只是普通的张量，但有一个非常特殊的附加功能：我们告诉 PyTorch 这些张量需要计算梯度。这使得 PyTorch 会记录在张量上执行的所有操作，以便在反向传播时自动计算梯度！

对于权重，我们在初始化之后设置 requires_grad，因为我们不希望该步骤包含在梯度中。（请注意，PyTorch 中尾随的 _ 表示该操作是就地执行的。）

我们在这里使用 Xavier 初始化来初始化权重（通过乘以 1/sqrt(n)）。

importmath

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 代码。

deflog_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

defmodel(xb):
    return log_softmax(xb @ weights + bias)

在上述内容中，@ 表示矩阵乘法操作。我们将在一批数据（本例中为 64 张图像）上调用我们的函数。这就是一次前向传播。请注意，由于我们是从随机权重开始的，因此在此阶段我们的预测不会比随机猜测更好。

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)

tensor([-2.5452, -2.0790, -2.1832, -2.6221, -2.3670, -2.3854, -2.9432, -2.4391,
        *1.8657, -2.0355], grad_fn=<SelectBackward0>) torch.Size([64, 10])

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

让我们实现负对数似然作为损失函数（同样，我们可以直接使用标准的 Python）：

defnll(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.4020, grad_fn=<NegBackward0>)

我们还需要实现一个函数来计算模型的准确率。对于每个预测，如果最大值的索引与目标值匹配，则预测正确。

defaccuracy(out, yb):
    preds = torch.argmax(out, dim=1)
    return (preds == yb).float().mean()

让我们检查随机模型的准确性，以便观察随着损失降低，准确性是否有所提高。

print(accuracy(preds, yb))

tensor(0.0938)

我们现在可以运行训练循环。在每次迭代中，我们将：

选择一个小批量数据（大小为 bs）
使用模型进行预测
计算损失
loss.backward() 更新模型的梯度，在本例中为 weights 和 bias。

我们现在使用这些梯度来更新权重和偏置。我们在 torch.no_grad() 上下文管理器中进行此操作，因为我们不希望这些操作被记录到下一次梯度计算中。您可以在此处了解更多关于 PyTorch 的 Autograd 如何记录操作的信息：这里。

然后我们将梯度设置为零，以便为下一次循环做好准备。否则，我们的梯度将记录所有已发生操作的累计值（即 loss.backward() 会将梯度累加到已存储的值上，而不是替换它们）。

您可以使用标准的 Python 调试器逐步调试 PyTorch 代码，从而检查每一步中各个变量的值。取消下面 set_trace() 的注释即可尝试。

fromIPython.core.debuggerimport 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_()

就这样：我们完全从零开始创建并训练了一个简单的神经网络（在这个例子中是一个逻辑回归模型，因为我们没有隐藏层）！

让我们检查一下损失和准确率，并与之前的结果进行比较。我们预计损失会减少，准确率会提高，而这些都已经实现了。

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

tensor(0.0813, grad_fn=<NegBackward0>) tensor(1.)

使用 `torch.nn.functional`

我们现在将重构代码，使其功能与之前相同，但我们会开始利用 PyTorch 的 nn 类来使代码更加简洁和灵活。从现在开始的每一步，我们都应该让代码变得更短、更易理解，和/或更灵活。

第一步也是最简单的一步，是通过用 torch.nn.functional 中的激活函数和损失函数（通常按照惯例导入到 F 命名空间中）来替换我们手写的函数，从而使代码更短。这个模块包含了 torch.nn 库中的所有函数（而库的其他部分则包含类）。除了广泛的损失函数和激活函数外，你还可以在这里找到一些用于创建神经网络的便捷函数，例如池化函数。（还有一些用于卷积、线性层等的函数，但正如我们将看到的，这些通常更适合使用库的其他部分来处理。）

如果您使用的是负对数似然损失和 log softmax 激活函数，那么 PyTorch 提供了一个名为 F.cross_entropy 的函数，它将两者结合在一起。因此，我们甚至可以从模型中移除激活函数。

importtorch.nn.functionalasF

loss_func = F.cross_entropy

defmodel(xb):
    return xb @ weights + bias

请注意，我们不再在 model 函数中调用 log_softmax。让我们确认一下损失和准确率与之前相同：

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

tensor(0.0813, grad_fn=<NllLossBackward0>) tensor(1.)

使用 `nn.Module` 进行重构

接下来，我们将使用 nn.Module 和 nn.Parameter 来构建一个更清晰、更简洁的训练循环。我们继承 nn.Module（它本身是一个类，并且能够跟踪状态）。在这种情况下，我们希望创建一个类来保存我们的权重、偏置以及前向传播的方法。nn.Module 拥有许多属性和方法（例如 .parameters() 和 .zero_grad()），我们将使用它们。

nn.Module（大写 M）是 PyTorch 特有的概念，它是一个我们会经常使用的类。不要将 nn.Module 与 Python 中的（小写 m）模块概念混淆，后者是一个可以导入的 Python 代码文件。

fromtorchimport nn

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

    defforward(self, xb):
        return xb @ self.weights + self.bias

由于我们现在使用的是一个对象而不是仅仅使用一个函数，我们首先需要实例化我们的模型：

model = Mnist_Logistic()

现在我们可以像之前一样计算损失。注意，nn.Module 对象可以像函数一样使用（即可调用），但在背后，PyTorch 会自动调用我们的 forward 方法。

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

tensor(2.3096, grad_fn=<NllLossBackward0>)

以前，在我们的训练循环中，我们必须通过名称逐个更新每个参数的值，并手动将每个参数的梯度单独清零，就像这样：

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 定义），使这些步骤更加简洁，并减少因忘记某些参数而导致的错误，尤其是在我们拥有更复杂模型的情况下：

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

我们将把训练循环封装在一个 fit 函数中，以便稍后可以再次运行它。

deffit():
    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.0821, grad_fn=<NllLossBackward0>)

使用 `nn.Linear` 进行重构

我们继续重构代码。与其手动定义和初始化 self.weights 和 self.bias，并计算 xb @ self.weights + self.bias，我们将使用 PyTorch 的 nn.Linear 类来实现线性层，它为我们完成了所有这些工作。PyTorch 提供了许多预定义的层类型，可以极大地简化我们的代码，并且通常还能使其运行得更快。

classMnist_Logistic(nn.Module):
    def__init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

    defforward(self, xb):
        return self.lin(xb)

我们实例化模型并按照之前的方式计算损失：

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

tensor(2.3313, grad_fn=<NllLossBackward0>)

我们仍然可以使用与之前相同的 fit 方法。

fit()

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

tensor(0.0819, grad_fn=<NllLossBackward0>)

使用 `torch.optim` 进行重构

PyTorch 还提供了一个包含各种优化算法的包，torch.optim。我们可以使用优化器中的 step 方法来执行前向步骤，而不是手动更新每个参数。

这将使我们能够替换之前手动编写的优化步骤：

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

并改为直接使用：

opt.step()
opt.zero_grad()

(optim.zero_grad() 将梯度重置为 0，我们需要在为下一个 minibatch 计算梯度之前调用它。)

fromtorchimport optim

我们将定义一个简单的函数来创建我们的模型和优化器，以便将来可以重复使用它。

defget_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.2659, grad_fn=<NllLossBackward0>)
tensor(0.0810, grad_fn=<NllLossBackward0>)

使用 Dataset 进行重构

PyTorch 提供了一个抽象的 Dataset 类。任何具有 __len__ 函数（由 Python 的标准 len 函数调用）和 __getitem__ 函数（用于索引）的对象都可以被视为 Dataset。本教程详细演示了如何创建一个自定义的 FacialLandmarkDataset 类作为 Dataset 的子类。

PyTorch 的 TensorDataset 是一个包装张量的 Dataset。通过定义长度和索引方式，它还能够让我们在张量的第一个维度上进行迭代、索引和切片操作。这将使我们在训练时更容易在同一行中访问自变量和因变量。

fromtorch.utils.dataimport TensorDataset

x_train 和 y_train 可以合并到一个 TensorDataset 中，这样更方便进行迭代和切片操作。

train_ds = TensorDataset(x_train, y_train)

之前，我们必须分别遍历 x 和 y 值的小批量数据：

xb = x_train[start_i:end_i]
yb = y_train[start_i:end_i]

现在，我们可以将这两个步骤合并在一起：

xb,yb = train_ds[i*bs : i*bs+bs]

model, 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.0826, grad_fn=<NllLossBackward0>)

使用 `DataLoader` 进行重构

PyTorch 的 DataLoader 负责管理批处理。您可以从任何 Dataset 创建 DataLoader。DataLoader 使得遍历批处理变得更加容易。相比于使用 train_ds[i*bs : i*bs+bs]，DataLoader 会自动为我们提供每个小批量的数据。

fromtorch.utils.dataimport DataLoader

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs)

之前，我们的循环是这样遍历批次 (xb, yb) 的：

for i in range((n-1)//bs + 1):
    xb,yb = train_ds[i*bs : i*bs+bs]
    pred = model(xb)

现在，我们的循环更加简洁，因为 (xb, yb) 会自动从数据加载器中加载：

for xb,yb in train_dl:
    pred = model(xb)

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.0818, grad_fn=<NllLossBackward0>)

得益于 PyTorch 的 nn.Module、nn.Parameter、Dataset 和 DataLoader，我们的训练循环现在显著简化且更易于理解。接下来，我们将尝试添加在实践中创建有效模型所需的基本功能。

添加验证

在第1节中，我们只是尝试为训练数据设置一个合理的训练循环。实际上，您始终还应该有一个验证集，以便确定是否出现过拟合。

打乱训练数据对于防止批次之间的相关性和过拟合是重要的。另一方面，无论我们是否打乱验证集，验证损失都是相同的。由于打乱操作需要额外的时间，因此没有必要对验证数据进行打乱。

我们将为验证集使用比训练集大两倍的批量大小。这是因为验证集不需要反向传播，因此占用的内存更少（它不需要存储梯度）。我们利用这一点来使用更大的批量大小，从而更快地计算损失。

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.3048)
1 tensor(0.2872)

创建 fit() 和 get_data()

我们现在将进行一点重构。由于我们在计算训练集和验证集的损失时经历了类似的过程两次，因此我们将其封装成一个独立的函数 loss_batch，用于计算一个批次的损失。

对于训练集，我们传入一个优化器，并使用它来执行反向传播。而对于验证集，我们不传入优化器，因此该方法不会执行反向传播。

defloss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)

fit 方法执行必要的操作来训练我们的模型，并计算每个训练周期的训练损失和验证损失。

importnumpyasnp

deffit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)

        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print(epoch, val_loss)

get_data 返回训练集和验证集的数据加载器。

defget_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )

现在，我们获取数据加载器并拟合模型的整个过程只需3行代码即可完成：

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
model, opt = get_model()
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.2939354367017746
1 0.3258970756947994

您可以使用这三行基础代码来训练多种模型。让我们看看是否能用它们来训练一个卷积神经网络（CNN）！

切换到卷积神经网络 (CNN)

我们现在将构建一个包含三个卷积层的神经网络。由于前一节中的函数没有对模型形式做出任何假设，因此我们无需修改即可使用它们来训练CNN。

我们将使用PyTorch预定义的Conv2d类作为我们的卷积层。我们定义了一个包含3个卷积层的CNN。每个卷积层后面都跟着一个ReLU激活函数。最后，我们执行平均池化操作。（注意，view是PyTorch中与Numpy的reshape功能对应的操作。）

classMnist_CNN(nn.Module):
    def__init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1)

    defforward(self, xb):
        xb = xb.view(-1, 1, 28, 28)
        xb = F.relu(self.conv1(xb))
        xb = F.relu(self.conv2(xb))
        xb = F.relu(self.conv3(xb))
        xb = F.avg_pool2d(xb, 4)
        return xb.view(-1, xb.size(1))

lr = 0.1

Momentum 是随机梯度下降的一种变体，它考虑了之前的更新，通常能加快训练速度。

model = Mnist_CNN()
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.3646130460739136
1 0.26228193019628526

使用 `nn.Sequential`

torch.nn 提供了另一个方便的类来简化我们的代码：Sequential。Sequential 对象会按顺序运行其包含的每个模块。这是一种更简单的编写神经网络的方式。

为了利用这一点，我们需要能够轻松地根据给定函数定义一个自定义层。例如，PyTorch 没有视图层，我们需要为我们的网络创建一个。Lambda 可以创建一个层，然后我们可以在使用 Sequential 定义网络时使用它。

classLambda(nn.Module):
    def__init__(self, func):
        super().__init__()
        self.func = func

    defforward(self, x):
        return self.func(x)


defpreprocess(x):
    return x.view(-1, 1, 28, 28)

使用 Sequential 创建的模型非常简单：

model = nn.Sequential(
    Lambda(preprocess),
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AvgPool2d(4),
    Lambda(lambda x: x.view(x.size(0), -1)),
)

opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.3330025281429291
1 0.22993727023601532

封装 `DataLoader`

我们的CNN相当简洁，但它仅适用于MNIST，原因如下：

它假设输入是一个 28*28 的长向量
它假设最终的 CNN 网格大小为 4*4（因为这是我们使用的平均池化核大小）

让我们消除这两个假设，以便我们的模型能够处理任何二维单通道图像。首先，我们可以通过将数据预处理移动到生成器中来移除初始的 Lambda 层：

defpreprocess(x, y):
    return x.view(-1, 1, 28, 28), y


classWrappedDataLoader:
    def__init__(self, dl, func):
        self.dl = dl
        self.func = func

    def__len__(self):
        return len(self.dl)

    def__iter__(self):
        for b in self.dl:
            yield (self.func(*b))

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

接下来，我们可以用 nn.AdaptiveAvgPool2d 替换 nn.AvgPool2d，这样我们就能定义所需的输出张量大小，而不是根据已有的输入张量大小。因此，我们的模型能够处理任意大小的输入。

model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    Lambda(lambda x: x.view(x.size(0), -1)),
)

opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

让我们来试一下：

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.3212135115623474
1 0.21439074140787123

使用您的加速器

如果您有幸能够使用 CUDA 等加速器（您可以从大多数云服务提供商那里以大约每小时 0.50 美元的价格租用），您可以使用它来加速代码。首先，请确保您的加速器在 PyTorch 中正常工作：

# If the current accelerator is available, we will use it. Otherwise, we use the CPU.
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")

Using cuda device

让我们更新 preprocess 方法，将批量数据移动到加速器上：

defpreprocess(x, y):
    return x.view(-1, 1, 28, 28).to(device), y.to(device)


train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

最后，我们可以将模型移至加速器。

model.to(device)
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

现在您会发现它运行得更快了：

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.1809617340326309
1 0.16968207444548608

总结

我们现在有了一个通用的数据管道和训练循环，您可以使用它来训练多种类型的 PyTorch 模型。为了了解现在训练模型可以变得多么简单，请查看 mnist_sample 笔记本。

当然，您可能还希望添加许多功能，例如数据增强、超参数调优、训练监控、迁移学习等。这些功能在 fastai 库中已经实现，该库采用了与本教程相同的设计方法，为希望进一步提升模型的从业者提供了自然的下一步。

在本教程开始时，我们承诺将通过示例解释 torch.nn、torch.optim、Dataset 和 DataLoader。现在，让我们总结一下我们所看到的内容：

torch.nn:

Module: 创建一个可调用的对象，其行为类似于函数，但也可以包含状态（例如神经网络的层权重）。它知道它包含哪些 Parameter（参数），并且可以清零所有参数的梯度，遍历它们以进行权重更新等。

Parameter: 对张量的封装，用于告知 Module 它包含需要在反向传播期间更新的权重。只有设置了 requires_grad 属性的张量才会被更新。

functional: 一个模块（通常按照惯例导入到 F 命名空间中），其中包含激活函数、损失函数等，以及无状态的层版本，如卷积层和线性层。

torch.optim: 包含优化器，例如 SGD，它们在反向传播步骤中更新 Parameter 的权重

Dataset: 一个包含 __len__ 和 __getitem__ 方法的抽象接口对象，包括 PyTorch 提供的类，如 TensorDataset

DataLoader: 接收任意 Dataset 并创建一个返回数据批次的迭代器。

下载 Python 源代码: nn_tutorial.py

下载 Jupyter 笔记本: nn_tutorial.ipynb

torch.nn 到底是什么？

MNIST 数据准备

从零开始构建神经网络（不使用 torch.nn）

使用 torch.nn.functional

使用 nn.Module 进行重构

使用 nn.Linear 进行重构

使用 torch.optim 进行重构