[人工智能]pytorch TORCH.NN 到底是什么？

PyTorch 提供了设计精美的模块和类torch.nn、 torch.optim、 Dataset和DataLoader 来帮助创建和训练神经网络。为了充分利用它们的力量并针对需求灵活的定制它们，需要真正了解它们在做什么。为了加深这种理解，我们将首先在 MNIST 数据集上训练基本的神经网络，而不使用这些模型的任何特性；最初将只使用PyTorch最基本的张量功能。然后，将逐步从torch.nn、torch.optim、Dataset或 DataLoader中添加一个功能，针对性的演示每个部分的功能，以及它如何使代码更简洁或更灵活。

本教程假设您已经安装了 PyTorch，并且熟悉张量操作的基础知识。 （如果熟悉 Numpy 数组操作，会发现和这里使用的 PyTorch 张量操作几乎相同）。

1. 设置 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/master/_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")
print(x_train.shape)

(50000, 784)

PyTorch 使用torch.tensor，而不是 numpy 数组，所以这里需要转换数据为tensor。

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)

2. 从头开始构建神经网络（不使用 torch.nn）

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

PyTorch 提供了用 随机数或零 创建并初始化张量的方法，这里使用这些方法为简单的线性模型创建权重和偏差。这些只是常规张量，还有一个非常特殊的标识：这个标识告诉 PyTorch 它们需要梯度。这会导致 PyTorch 记录在张量上完成的所有操作，以便它可以在反向传播期间自动计算梯度！

对于权重，在初始化之后设置requires_grad标识，因为不希望该步骤包含在梯度中。（请注意，_PyTorch 中的尾随表示该操作是in-place执行的。）

注意事项：

在这里使用 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 函数（或callable对象）作为模型！所以让我们编写一个普通的矩阵乘法和广播加法来创建一个简单的线性模型。这里还需要一个激活函数，所以需要编写并使用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 张图像）调用我们的函数。这是一个forward，即前向传播操作。注意，在这个阶段预测的不会比随机更好，因为用随机权重开始做的初始化操作。

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.5078, -2.1811, -2.3981, -2.1644, -2.4276, -1.7887, -2.7199, -2.2655,
        -2.4138, -2.4608], grad_fn=<SelectBackward0>) torch.Size([64, 10])

如上所见，preds张量不仅包含张量值，还包含梯度函数（grad_fn）。稍后将使用它来进行反向传播。

下边实现负对数似然作为损失函数（同样，可以只使用标准 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.3068, grad_fn=<NegBackward0>)

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

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

检查一下随机模型的准确性，这样就可以看到我们的准确性是否随着损失的提高而提高。

print(accuracy(preds, yb))

tensor(0.1562)

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

• 选择一小批数据（大小为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.0812, grad_fn=<NegBackward0>) tensor(1.)

3. 使用 torch.nn.functional

现在将重构代码，使其与之前的代码完成同样的功能，只是将开始利用 PyTorch 的nn类来使其更加简洁和灵活。从这里开始的每一步，都会改进代码：更短、更易理解和或更灵活。

第一步也是最简单的一步是通过将手写的激活和损失函数替换为 from 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

注意，不再调用log_softmax该model函数。和之前一样让再确认一下损失和准确率：

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

tensor(0.0812, grad_fn=<NllLossBackward0>) tensor(1.)

4. 使用 nn.Module 重构模型

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

注意事项：

nn.Module（大写 M）是 PyTorch 特定的概念，也是经常使用的类。nn.Module不要与Python （小写）module 概念混淆，module是可以导入的 Python 代码文件。m

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

以前对于训练循环，必须按名称更新每个参数的值，并分别手动将每个参数的 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 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.0819, grad_fn=<NllLossBackward0>)

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

tensor(2.3217, grad_fn=<NllLossBackward0>)

仍然可以使用与fit以前相同的方法。

fit()

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

tensor(0.0822, grad_fn=<NllLossBackward0>)

6. 重构优化算法

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.2914, grad_fn=<NllLossBackward0>)
tensor(0.0822, grad_fn=<NllLossBackward0>)

7. 使用Dataset 重构代码

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

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

8. 使用 DataLoader 重构

PytorchDataLoader负责管理批次。可以创建DataLoader从任何一个Dataset. 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.0808, grad_fn=<NllLossBackward0>)

多亏了 Pytorch 的nn.Module、nn.Parameter、Dataset和DataLoader，训练循环现在代码量明显更少并且更容易理解。现在尝试添加在实践中创建有效模型所需的基本功能。

9. 添加验证

在第 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.3142)
1 tensor(0.2956)

10. 创建 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),
    )

现在，获取数据加载器和拟合模型的整个过程可以在 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.32458346838951113
1 0.29746281175613404

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

11. 切换到 CNN

现在要构建具有三层卷积层的神经网络。因为上一节中的任何函数都没有假设模型形式，所以无需任何修改就可以使用它们来训练 CNN。

将使用 Pytorch 的预定义 Conv2d类作为卷积层。定义了一个具有 3 个卷积层的 CNN。每个卷积后面跟着一个 ReLU。最后，执行平均池化。（请注意，这view是 PyTorch 的 numpy’s 版本 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.35496450316905975
1 0.2666731256365776

12. nn.Sequential

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

为了利用这一点，需要能够从给任意函数轻松定义自定义层。 例如，PyTorch 没有视图层，需要为网络创建一个视图层。Lambda 将创建一个层，然后可以在使用Sequential定义网络时使用该Lambda层 。

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.34979644216299055
1 0.23173938195705412

13. 包装数据加载器

• 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):
        batches = iter(self.dl)
        for b in batches:
            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.3121746739387512
1 0.22073891738653184

14. 使用 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.2050167937517166
1 0.1815803111433983

15. 总结和思考

现在有一个通用的数据管道和训练循环，可以使用它来使用 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.

原文地址