• 教程 >
  • 什么是 torch.nn 真的
快捷方式

什么是 torch.nn 真的

创建于:2018 年 12 月 26 日 | 最后更新:2025 年 1 月 24 日 | 最后验证:2024 年 11 月 05 日

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

我们建议将本教程作为 notebook 而不是脚本运行。要下载 notebook (.ipynb) 文件,请点击页面顶部的链接。

PyTorch 提供了优雅设计的模块和类 torch.nntorch.optimDatasetDataLoader,以帮助您创建和训练神经网络。为了充分利用它们的功能并根据您的问题进行自定义,您需要真正理解它们究竟在做什么。为了培养这种理解,我们将首先在 MNIST 数据集上训练基本的神经网络,而不使用这些模型中的任何功能;我们最初将仅使用最基本的 PyTorch 张量功能。然后,我们将一次递增地添加 torch.nntorch.optimDatasetDataLoader 中的一个功能,准确展示每个部分的作用,以及它如何使代码更简洁或更灵活。

本教程假定您已经安装了 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)
nn tutorial
(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 甚至会自动为您的函数创建快速加速器或矢量化 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() 更新模型的梯度,在本例中为 weightsbias

我们现在使用这些梯度来更新权重和偏差。我们在 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 库中的所有函数(而库的其他部分包含类)。除了各种各样的损失函数和激活函数外,您还可以在此处找到一些用于创建神经网络的便捷函数,例如池化函数。(还有用于进行卷积、线性层等的函数,但正如我们将看到的,这些通常最好使用库的其他部分来处理。)

如果您使用负对数似然损失和 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))
tensor(0.0813, grad_fn=<NllLossBackward0>) tensor(1.)

使用 nn.Module 重构

接下来,我们将使用 nn.Modulenn.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

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

现在我们可以像以前一样计算损失。请注意,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 重构

我们继续重构我们的代码。我们将使用 Pytorch 类 nn.Linear 用于线性层,而不是手动定义和初始化 self.weightsself.bias,并计算 xb  @ self.weights + self.bias,它为我们完成了所有这些工作。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()

而是只使用

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

使用 Dataset 重构

PyTorch 有一个抽象的 Dataset 类。Dataset 可以是任何具有 __len__ 函数(由 Python 标准 len 函数调用)和 __getitem__ 函数作为索引到其中的方法。 本教程 逐步介绍了如何创建自定义 FacialLandmarkDataset 类作为 Dataset 的子类。

PyTorch 的 TensorDataset 是一个包装张量的 Dataset。通过定义长度和索引方式,这也为我们提供了一种沿张量的第一个维度进行迭代、索引和切片的方法。这将使我们更容易在同一行中访问自变量和因变量,就像我们训练一样。

from torch.utils.data import TensorDataset

x_trainy_train 都可以组合在单个 TensorDataset 中,这将更容易迭代和切片。

以前,我们必须分别迭代 xy 值的小批量

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 创建 DataLoaderDataLoader 使迭代批次变得更容易。而不是必须使用 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))
tensor(0.0818, grad_fn=<NllLossBackward0>)

感谢 PyTorch 的 nn.Modulenn.ParameterDatasetDataLoader,我们的训练循环现在已大大缩小并且更易于理解。现在,让我们尝试添加在实践中创建有效模型所需的基本功能。

添加验证

在第 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.BatchNorm2dnn.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 运行必要的 operations 来训练我们的模型,并计算每个 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

你可以使用这基本的 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

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 没有 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

包装 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.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

使用你的 Accelerator

如果你足够幸运能够访问像 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 以将批次移动到加速器

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)
0 0.1805140619158745
1 0.17031861956119537

结束语

我们现在有了一个通用的数据管道和训练循环,你可以使用它来使用 Pytorch 训练多种类型的模型。要了解现在训练模型可以多么简单,请查看 mnist_sample notebook

当然,你还需要添加许多东西,例如数据增强、超参数调整、监控训练、迁移学习等等。这些功能在 fastai 库中可用,该库是使用本教程中展示的相同设计方法开发的,为希望进一步发展其模型的从业者提供了自然的下一步。

我们在本教程开始时承诺,我们将通过示例解释 torch.nntorch.optimDatasetDataLoader 中的每一个。所以让我们总结一下我们所看到的

  • torch.nn:

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

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

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

  • torch.optim:包含优化器,例如 SGD,它在反向步骤期间更新 Parameter 的权重

  • Dataset:具有 __len____getitem__ 的对象的抽象接口,包括 Pytorch 提供的类,例如 TensorDataset

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

脚本总运行时间: ( 0 分钟 36.523 秒)

由 Sphinx-Gallery 生成的图库

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获取你的问题解答

查看资源