Shortcuts

什么是`torch.nn` 真正是

Created On: Dec 26, 2018 | Last Updated: Jan 24, 2025 | Last Verified: Nov 05, 2024

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

我们推荐以笔记本而不是脚本方式运行此教程。要下载笔记本(.ipynb)文件,请点击页面顶部的链接。

PyTorch提供了设计优雅的模块和类:torch.nntorch.optimDatasetDataLoader 来帮助您创建和训练神经网络。为了充分利用它们的能力并为您的问题定制它们,您需要真正了解它们的具体用途。为了发展这种理解,我们首先将训练一个基本的神经网络来处理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数组格式存在,并使用Python特定的序列化格式Pickle进行存储。

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)
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 <https://course.fast.ai>`_进行学习)。

PyTorch提供了创建随机或零填充张量的方法,我们将使用它来创建简单线性模型的权重和偏置。这些只是普通张量,但有一个特别的功能:我们告诉PyTorch它们需要梯度。这使得PyTorch能够记录张量上所做的所有操作,从而在反向传播时*自动*计算梯度!

对于权重,我们在初始化之后设置``requires_grad``,因为我们不希望这个步骤也被包含在梯度计算中。(注意PyTorch中带有后缀``_``的操作表示在原地修改)。

备注

我们在这里使用`Xavier初始化 <http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf>`_(通过乘以``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.5906, -2.6994, -2.0356, -1.9248, -2.5879, -2.2701, -2.3471, -2.4051,
        -1.8955, -2.7084], 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.3025, grad_fn=<NegBackward0>)

我们还可以实现一个函数来计算模型的准确率。对于每个预测,如果最大值的索引与目标值匹配,那么预测就是正确的。

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

我们检查一下随机模型的准确率,以便在损失改善时看看准确率是否有提高。

print(accuracy(preds, yb))
tensor(0.0625)

我们现在可以运行一个训练循环。对于每次迭代,我们将:

  • 选择一个数据小批量(大小``bs``)

  • 使用模型进行预测

  • 计算损失

  • loss.backward()``会更新模型的梯度,在本例中是``weights``和``bias

我们现在利用这些梯度更新权重和偏置。我们在``torch.no_grad()``上下文管理器中完成此操作,因为我们不希望这些动作被记录到下一次梯度计算中。可以在这里了解更多有关PyTorch自动梯度记录的信息:《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.0830, 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.0830, grad_fn=<NllLossBackward0>) tensor(1.)

使用``nn.Module``进行重构

接下来,我们使用``nn.Module``和``nn.Parameter``实现更清晰、更简洁的训练循环。我们继承``nn.Module``(它本身是一个类,能够跟踪状态)。在本例中,我们想创建一个包含权重、偏置和前向步骤方法的类。nn.Module``具有许多属性和方法(例如.parameters()``和``.zero_grad()``),我们稍后会使用。

备注

``nn.Module``(大写的M)是PyTorch特定的概念,是我们将经常使用的一个类。需要注意的是,它与Python中的`模块 <https://docs.python.org/3/tutorial/modules.html>`_(小写的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.3390, 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.0847, grad_fn=<NllLossBackward0>)

使用``nn.Linear``进行重构

我们继续重构代码。我们不再手动定义和初始化``self.weights``和``self.bias``,并计算``xb @ self.weights + self.bias``。而是使用PyTorch类`nn.Linear <https://pytorch.org/docs/stable/nn.html#linear-layers>`_来创建线性层,这将为我们完成所有这些工作。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.2535, grad_fn=<NllLossBackward0>)

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

fit()

print(loss_func(model(xb), yb))
tensor(0.0811, 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.3887, grad_fn=<NllLossBackward0>)
tensor(0.0818, grad_fn=<NllLossBackward0>)

使用Dataset进行重构

PyTorch有一个抽象的`Dataset`类。任何具有``__len__``函数(由Python标准函数``len``调用)和``__getitem__``函数(作为索引方式)的东西都可以是Dataset类。`这个教程 <https://pytorch.org/tutorials/beginner/data_loading_tutorial.html>`_展示了一个通过创建一个定制的``FacialLandmarkDataset``类作为``Dataset``子类的例子。

PyTorch的`TensorDataset <https://pytorch.org/docs/stable/_modules/torch/utils/data/dataset.html#TensorDataset>`_是一个封装张量的Dataset类。通过定义长度和索引方式,这也为我们提供了一种迭代、索引和切片张量第一个维度的方式。这将使在训练中同时访问自变量和因变量变得更容易。

from torch.utils.data import TensorDataset

``x_train``和``y_train``可以组合到一个``TensorDataset``中,这会让迭代和切片更简便。

之前,我们必须分别迭代``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.0811, grad_fn=<NllLossBackward0>)

使用``DataLoader``进行重构

PyTorch的``DataLoader``负责管理批次。你可以从任何Dataset中创建一个``DataLoader``。``DataLoader``使迭代批次更加方便,而不用像之前使用``train_ds[i*bs : i*bs+bs]``的方式,它会自动给我们每个小批量。

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

得益于 PyTorch 的 nn.Modulenn.ParameterDatasetDataLoader,我们的训练循环现在大大缩短且更易于理解。现在,让我们尝试添加一些实际中创建高效模型所需的基本功能。

添加验证

在第一部分中,我们仅仅是在为训练数据设置一个合理的训练循环。实际上,你始终应该有一个`验证集 <https://www.fast.ai/2017/11/13/validation-sets/>`_,以识别是否发生了过拟合。

对训练数据进行`随机化 <https://www.quora.com/Does-the-order-of-training-data-matter-when-training-neural-networks>`_ 非常重要,可以防止批量间的相关性并避免过拟合。另一方面,无论是否对验证集进行随机化,验证损失都会是相同的。由于随机化需要额外的时间,因此随机化验证数据没有意义。

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

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)

我们将在每个周期结束时计算并打印验证损失。

(注意,我们在训练前始终调用 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.3108)
1 tensor(0.3200)

创建 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 执行必要的操作来训练我们的模型,并为每个周期计算训练和验证损失。

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

现在,我们获取数据加载器和拟合模型的过程可以用三行代码完成:

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.29459090477228167
1 0.3131585433602333

你可以使用这三行代码来训练各种类型的模型。让我们看看是否可以用它们来训练卷积神经网络 (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)
0 0.5929948786735535
1 0.5213978308200836

使用 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.3507859182596207
1 0.2347101690530777

封装 DataLoader

我们的 CNN 代码相当简洁,但它只能在 MNIST 数据集上运行,因为:
  • 它假设输入是一个 28*28 的长向量

  • 它假设最终的 CNN 网格尺寸是 4*4(因为我们使用了 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.3828958973646164
1 0.25244341576099394

使用你的 加速器

如果你足够幸运拥有类似 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.20608520213365555
1 0.20726044411063194

结束语

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

当然,仍有许多你可能希望添加的功能,例如数据增强、超参数调优、训练监控、迁移学习等。这些功能在 fastai 库中提供,该库使用与本教程中展示的设计方法相同的方式开发,对于希望进一步提升模型的实践者来说是自然的下一步。

我们在本教程的开头承诺将通过示例讲解 torch.nntorch.optimDatasetDataLoader。现在,让我们总结一下我们所看到的内容:

  • torch.nn:

    • Module: 创建一个既可以像函数一样调用又可以包含状态(如神经网络层权重)的实例。它知道它包含哪些 ``Parameter``(参数),并可以清零它们的梯度、循环访问它们以进行权重更新等。

    • Parameter: Tensor 的一个封装,告知 Module 它拥有需要在反向传播期间更新的权重。只有设置了 requires_grad 属性的张量才会被更新。

    • functional: 一个模块(通常根据惯例导入到 F 命名空间中),包含激活函数、损失函数等内容,还包括无状态版本的层,如卷积层和线性层。

  • torch.optim: 包含诸如 SGD 的优化器,用于在反向步骤中更新 Parameter 的权重。

  • Dataset: 一个抽象接口,具有 __len____getitem__ 的对象,包括 PyTorch 提供的 TensorDataset 等类。

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

脚本总运行时间: (1 分钟 7.140 秒)

画廊由 Sphinx-Gallery 生成

文档

访问 PyTorch 的详细开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

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

查看资源