torch.nn 到底是什么？
原文：
pytorch.org/tutorials/beginner/nn_tutorial.html译者：飞龙
注意
点击这里下载完整示例代码
作者： 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)
(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进行重构
我们继续重构我们的代码。我们将使用 Pytorch 类nn.Linear来代替手动定义和初始化self.weights和self.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()
而是使用：
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>)
使用 Dataset 进行重构
PyTorch 有一个抽象的 Dataset 类。一个 Dataset 可以是任何具有__len__函数（由 Python 的标准len函数调用）和__getitem__函数作为索引方式的东西。这个教程演示了创建一个自定义的FacialLandmarkDataset类作为Dataset子类的一个很好的例子。
PyTorch 的TensorDataset是一个包装张量的数据集。通过定义长度和索引方式，这也为我们提供了一种在张量的第一维度上进行迭代、索引和切片的方式。这将使我们更容易在训练时同时访问独立变量和因变量。
from torch.utils.data import 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会自动给我们每个小批次。
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，用于计算一个 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
您可以使用这基本的 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)
0 0.35247018008232117 1 0.25782823679447175
使用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)
0 0.3226209937572479 1 0.2234949318766594
包装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.3148617018699646 1 0.20678156037330628
使用您的 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.17924857176542283 1 0.17124842552542688
PyTorch 2.2 中文官方教程（四）（2）https://developer.aliyun.com/article/1482493
