torch.nn 究竟是什么？

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：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 下载数据集。我们将只在需要时导入模块，这样您可以清楚地看到每个点使用了什么。

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 数组格式存储，并使用 pickle 存储，pickle 是 Python 特有的数据序列化格式。

import pickle
import 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）。让我们看一下其中一个；我们需要先将其重塑为 2d。

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)

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

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

从零开始构建神经网络（不包含 torch.nn ）¶

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

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

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

备注

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

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

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 x
preds = model(xb)  # predictions
preds[0], preds.shape
print(preds[0], preds.shape)

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

让我们再次实现负对数似然作为损失函数（我们仍然可以使用标准的 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))

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

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

让我们检查随机模型的准确率，以便我们可以在损失改进时看到准确率是否有所提高。

print(accuracy(preds, yb))

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

• 选择一个大小为 bs 的小批量数据

• 使用模型进行预测

• 计算损失

• 更新模型的梯度，在这种情况下， weights 和 bias 。

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

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

提示

您可以使用标准的 Python 调试器逐步执行 PyTorch 代码，这样您就可以在每一步检查各种变量的值。取消注释下面的 set_trace() 以尝试它。

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

就这样：我们从零开始创建并训练了一个最小的神经网络（在这种情况下，是一个逻辑回归，因为我们没有隐藏层）！

让我们检查损失和准确率，并将它们与我们之前得到的结果进行比较。我们预计损失会下降，准确率会上升，它们确实如此。

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

使用 torch.nn.functional ¶

现在，我们将重构我们的代码，使其与之前做同样的事情，只是我们会开始利用 PyTorch 的 nn 类来使其更加简洁和灵活。从这一步开始，我们应该使我们的代码变得更短、更易于理解、更灵活。

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

如果您使用负对数似然损失和 log softmax 激活，Pytorch 提供了一个组合这两个功能的单一函数 F.cross_entropy 。因此，我们甚至可以移除模型中的激活函数。

import torch.nn.functional as F

loss_func = F.cross_entropy

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

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

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

使用 nn.Module 进行重构

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

备注

nn.Module （大写 M）是 PyTorch 特有的概念，是我们将大量使用的类。 nn.Module 不要与 Python 概念中的（小写 m ）模块混淆，模块是一个 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.bias

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

model = Mnist_Logistic()

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

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

之前在我们的训练循环中，我们必须按名称更新每个参数的值，并分别手动将每个参数的 grads 置零，就像这样：

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

现在我们可以利用 model.parameters()和 model.zero_grad()（这两个都由 PyTorch 定义）来使这些步骤更加简洁，并减少忘记某些参数的错误，尤其是如果我们有一个更复杂的模型的话：

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

我们将把我们的训练循环包装在一个 fit 函数中，这样我们就可以稍后再次运行它。

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

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

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

fit()

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

使用 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 的梯度之前需要调用它。）

from torch import optim

我们将定义一个小函数来创建我们的模型和优化器，这样我们就可以在将来重复使用它。

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

使用数据集重构

PyTorch 有一个抽象的 Dataset 类。数据集可以是任何具有 __len__ 函数（由 Python 的标准 len 函数调用）和 __getitem__ 函数作为索引方式的东西。本教程将带您通过创建一个自定义 FacialLandmarkDataset 类的示例，它是 Dataset 的子类。

PyTorch 的 TensorDataset 是一个封装张量的 Dataset。通过定义长度和索引方式，这也为我们提供了迭代、索引和切片张量第一维的方法。这将使我们更容易在训练过程中在同一行中访问独立变量和依赖变量。

from torch.utils.data import TensorDataset

x_train 和 y_train 可以合并为单个 TensorDataset ，这将更容易迭代和切片。

train_ds = TensorDataset(x_train, y_train)

之前，我们必须分别迭代 x 和 y 的 minibatch 值：

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

使用 DataLoader 进行重构

PyTorch 的 DataLoader 负责管理批次。您可以从任何 Dataset 创建一个 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)

之前，我们的循环是这样遍历批次的 (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))

多亏了 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))

创建 fit()和 get_data()方法

我们现在将进行一点自己的重构。由于我们两次都要计算训练集和验证集的损失，所以让我们将其制作成一个自己的函数，即 loss_batch ，该函数计算一个批次的损失。

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

def loss_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 运行必要的操作来训练我们的模型，并计算每个 epoch 的训练和验证损失。

import numpy as np

def fit(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 返回训练集和验证集的 dataloader。

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

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

切换到 CNN 模式

我们现在将使用三个卷积层来构建我们的神经网络。由于上一节中的所有函数都没有对模型形式做出任何假设，我们将能够直接使用它们来训练 CNN，无需任何修改。

我们将使用 PyTorch 预定义的 Conv2d 类作为我们的卷积层。我们定义了一个包含 3 个卷积层的 CNN。每个卷积后都跟一个 ReLU。最后，我们进行平均池化。（注意， view 是 PyTorch 对 Numpy 的 reshape 的版本）

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

    def forward(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

动量是随机梯度下降的一种变体，它也考虑了之前的更新，并且通常会导致训练速度更快。

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

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

使用 nn.Sequential ¶

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

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

class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


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

包装 DataLoader ¶

我们的 CNN 相当简洁，但它只适用于 MNIST，因为：

• 它假设输入是一个 28*28 的长向量

• 它假设最终的 CNN 网格大小是 4*4（因为这是我们使用的平均池化核大小）

让我们摆脱这两个假设，这样我们的模型就可以处理任何 2d 单通道图像。首先，我们可以通过将数据预处理移动到生成器中来移除初始的 Lambda 层：

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


class WrappedDataLoader:
    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.AvgPool2d 替换为 nn.AdaptiveAvgPool2d ，这样我们就可以定义我们想要的输出张量的大小，而不是我们拥有的输入张量的大小。因此，我们的模型将可以处理任何大小的输入。

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)

使用您的加速器 ¶

如果您有幸能访问到 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")

让我们更新 preprocess 将批次移动到加速器：

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

结束语 ¶

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

当然，您可能想要添加许多其他功能，例如数据增强、超参数调整、训练监控、迁移学习等等。这些功能在 fastai 库中都有提供，该库采用了与本文教程相同的设计方法，为希望将模型进一步发展的实践者提供了一个自然的下一步。

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

• torch.nn:

  • Module ：创建一个类似于函数的可调用对象，但也可以包含状态（例如神经网络层的权重）。它知道它包含的 Parameter （s），可以将其所有梯度置零，循环遍历它们以更新权重等。

  • Parameter ：一个用于张量的包装器，告诉 Module 它在反向传播期间需要更新权重。只有设置了 requires_grad 属性的张量才会被更新。

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

• torch.optim ：包含诸如 SGD 之类的优化器，这些优化器在反向传播步骤中更新 Parameter 的权重。

• Dataset ：具有 __len__ 和 __getitem__ 的抽象接口的对象，包括 PyTorch 提供的类 TensorDataset 等。

• DataLoader ：接受任何 Dataset 并创建一个迭代器，该迭代器返回数据批次。