什么是 torch.nn 实际上?

创建日期: 2018年12月26日 | 最后更新日期: 2024年12月3日 | 最后验证日期: 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，这是一种特定于 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）的一维扁平化数组。让我们看一下其中一张；首先需要将其重塑为二维形式。

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)

(50000, 784)

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

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

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 甚至会为您自动创建快速 GPU 或向量化 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)

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）：

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

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

def accuracy(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() 下面的代码以尝试一下。

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

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

使用 torch.nn.functional

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

第一步也是最简单的步骤是通过使用torch.nn.functional中的激活函数和损失函数来使我们的代码更短（这些函数通常被导入到命名空间F中）。这个模块包含了torch.nn库中的所有函数（而库中的其他部分则包含类）。除了广泛的损失和激活函数外，您还会在这里找到一些方便的创建神经网络的功能，例如池化函数。（库中也有一些进行卷积、线性层等操作的函数，但正如我们将看到的，这些通常最好使用库中的其他部分来处理。）

如果您使用的是负对数似然损失和对数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))

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

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

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 函数中，以便以后再次运行。

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

使用 nn.Linear 重构

我们继续重构我们的代码。Instead of manually defining and initializing self.weights and self.bias, and calculating 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.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，并且在计算下一个小批量的梯度之前需要调用此函数。)

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

tensor(2.2659, grad_fn=<NllLossBackward0>)
tensor(0.0810, grad_fn=<NllLossBackward0>)

Refactor using Dataset

PyTorch 有一个抽象的 Dataset 类。Dataset 可以是任何具有 __len__ 函数（由 Python 的标准 len 函数调用）和 __getitem__ 函数作为索引方式的对象。 本教程 介绍了如何创建一个自定义的 FacialLandmarkDataset 类，作为 Dataset 的子类。

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

from torch.utils.data import TensorDataset

Both x_train and 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 会自动为我们提供每个 minibatch。

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

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，该函数用于计算一个批次的损失。

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

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 返回用于训练和验证集的数据加载器。

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)

0 0.2939354367017746
1 0.3258970756947994

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

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

0 0.3646130460739136
1 0.26228193019628526

使用 nn.Sequential

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

为了充分利用这一点，我们需要能够轻松地从给定函数定义一个 自定义层。例如，PyTorch 没有 view 层，我们需要为此网络创建一个。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)

0 0.3330025281429291
1 0.22993727023601532

Wrapping DataLoader

Our CNN is fairly concise, but it only works with MNIST, because:

• 它假设输入是一个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)

0 0.3212135115623474
1 0.21439074140787123

使用您的GPU

如果你有幸能够使用一个支持CUDA的GPU（大多数云服务提供商可以提供租赁服务，价格约为每小时0.50美元），你可以利用它来加快代码运行速度。首先检查你的GPU是否能在Pytorch中正常工作：

print(torch.cuda.is_available())

True

然后为其创建一个设备对象：

dev = torch.device(
    "cuda") if torch.cuda.is_available() else torch.device("cpu")

让我们更新 preprocess 将批次移动到GPU：

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


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

最后，我们可以将模型移动到GPU。

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

你现在应该会发现运行得更快了：

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

0 0.18173173379898072
1 0.1716031593978405

Closing thoughts

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

当然，您还需要添加许多其他功能，例如数据增强、超参数调整、监控训练、迁移学习等等。fastai 库提供了这些功能，其开发采用了与本教程中所示相同的设计方法，为希望进一步提升模型的实践者提供了自然的下一步选择。

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

• torch.nn:

  • Module: creates a callable which behaves like a function, but can also contain state(such as neural net layer weights). It knows what Parameter (s) it contains and can zero all their gradients, loop through them for weight updates, etc.

  • Parameter: a wrapper for a tensor that tells a Module that it has weights that need updating during backprop. Only tensors with the requires_grad attribute set are updated

  • functional: a module(usually imported into the F namespace by convention) which contains activation functions, loss functions, etc, as well as non-stateful versions of layers such as convolutional and linear layers.

• torch.optim: Contains optimizers such as SGD, which update the weights of Parameter during the backward step

• Dataset: An abstract interface of objects with a __len__ and a __getitem__, including classes provided with Pytorch such as TensorDataset

• DataLoader: Takes any Dataset and creates an iterator which returns batches of data.