WHAT IS TORCH.NN REALLY?

作者：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 data setup

我们将使用经典的MNIST数据集，其中包括手绘数字的黑白图像（介于0和9之间）。我们将使用pathlib来处理路径（Python 3标准库的一部分），并将使用请求下载数据集。 我们只会在使用模块时导入模块，因此您可以准确地看到每个模块的使用情况。

from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"

PATH.mkdir(parents=True, exist_ok=True)

URL = "http://deeplearning.net/data/mnist/"
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)

Out:

(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
x_train, x_train.shape, y_train.min(), y_train.max()
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

Out:

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)

Neural net from scratch (no torch.nn)

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

PyTorch提供了创建随机或零填充张量的方法，我们将使用这些方法为简单的线性模型创建权重和偏差。 这些只是常规张量，有一个非常特别的补充：我们告诉PyTorch它们需要一个梯度。 这会导致PyTorch记录在张量上完成的所有操作，以便它可以在反向传播过程中自动计算梯度！

对于权重，我们在初始化之后设置requires_grad，因为我们不希望梯度中包含该步骤。 （请注意，PyTorch中的trailling _表示操作是就地执行的。）

我们在这里用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个图像）。 这是一个正向pass。 请注意，我们的预测在这个阶段不会比随机更好，因为我们从随机权重开始。

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)

Out:

tensor([-2.2588, -2.8888, -2.2290, -2.1785, -2.1712, -2.3731, -2.1631, -2.2584,
        -2.4279, -2.2624], grad_fn=<SelectBackward>) torch.Size([64, 10])

如您所见，preds张量不仅包含张量值，还包含梯度函数。 我们稍后会用它来做backprop。

让我们实现负对数似然用作损失函数（同样，我们可以使用标准Python）：

def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll

让我们用我们的随机模型检查我们的损失，这样我们就可以看到我们是否在一个backprop传递之后改进了。

yb = y_train[0:bs]
print(loss_func(preds, yb))

Out:

tensor(2.3516, grad_fn=<NegBackward>)

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

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

让我们检查一下随机模型的准确性，这样我们就可以看出随着损失的改善我们的准确度是否会提高。

print(accuracy(preds, yb))

Out:

tensor(0.0156)

我们现在可以运行一个训练循环。 对于每次迭代，我们将：选择一小批数据（大小为bs）

使用该模型进行预测

计算损失

loss.backward（）更新模型的梯度，在这种情况下，权重和偏差。

我们现在使用这些梯度来更新权重和偏差。 我们在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))

Out:

tensor(0.0803, grad_fn=<NegBackward>) tensor(1.)

Using 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

请注意，我们不再在模型函数中调用log_softmax。 让我们确认我们的损失和准确性与以前相同：

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

Out:

tensor(0.0803, grad_fn=<NllLossBackward>) tensor(1.)

Refactor using nn.Module

接下来，我们将使用nn.Module和nn.Parameter，以获得更清晰，更简洁的训练循环。 我们子类nn.Module（它本身是一个类，能够跟踪状态）。 在这种情况下，我们想要创建一个包含前向步骤的权重，偏差和方法的类。 nn.Module有许多我们将使用的属性和方法（例如.parameters（）和.zero_grad（））。

nn.Module（大写M）是PyTorch特定的概念，是我们将要使用的一个类。 nn.Module不要与（小写m）模块的Python概念混淆，后者是可以导入的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))

Out:

tensor(2.2681, grad_fn=<NllLossBackward>)

以前，对于我们的训练循环，我们必须按名称更新每个参数的值，并分别手动将每个参数的梯度归零，如下所示：

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

我们将我们的小训练循环包装在拟合函数中，以便我们稍后再次运行它。

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

Out:

tensor(0.0801, grad_fn=<NllLossBackward>)

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

Out:

tensor(2.3232, grad_fn=<NllLossBackward>)

我们仍然可以使用与以前相同的拟合方法。

fit()

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

Out:

tensor(0.0829, grad_fn=<NllLossBackward>)

Refactor using optim

Pytorch还有一个包含各种优化算法的软件包torch.optim。 我们可以使用优化器中的步骤方法来执行前进步骤，而不是手动更新每个参数。

这将让我们取代之前的手动编码优化步骤：

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

Out:

tensor(2.3002, grad_fn=<NllLossBackward>)
tensor(0.0801, grad_fn=<NllLossBackward>)

Refactor using Dataset

PyTorch有一个抽象的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))

Out:

tensor(0.0821, grad_fn=<NllLossBackward>)

Refactor using DataLoader

Pytorch的DataLoader负责管理批次。 您可以从任何数据集创建DataLoader。 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))

Out:

tensor(0.0813, grad_fn=<NllLossBackward>)

感谢Pytorch的nn.Module，nn.Parameter，Dataset和DataLoader，我们的训练循环现在变得更小，更容易理解。 现在让我们尝试添加在实践中创建有效模型所需的基本功能。

Add validation

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

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

（注意，我们总是在训练之前调用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))

Out:

0 tensor(0.3340)
1 tensor(0.3016)

Create fit() and get_data()

我们现在将对自己进行一些重构。 由于我们经历了两次计算训练集和验证集损失的类似过程，让我们将其转换为自己的函数loss_batch，它计算一个批次的损失。

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

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)

Out:

0 0.5046828230142594
1 0.3456133277416229

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

Switch to CNN

我们现在要用三个卷积层构建我们的神经网络。 因为上一节中没有任何函数假定模型形式的任何内容，我们将能够使用它们来训练CNN而无需任何修改。

我们将使用Pytorch的预定义Conv2d类作为卷积层。 我们定义了一个带有3个卷积层的CNN。 每个卷积后面都有一个ReLU。 最后，我们执行平均池化。 （注意，view是PyTorch的numpy shape版本）

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)

Out:

0 0.3969319754123688
1 0.2569236306667328

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)

Out:

0 0.35227406759262087
1 0.26811458735466004

Wrapping DataLoader

我们的CNN相当简洁，但它只适用于MNIST，因为：

• 它假设输入是28 * 28长矢量
• 它假设最终CNN网格大小为4 * 4（因为这是我们使用的平均池化核大小）

让我们摆脱这两个假设，因此我们的模型适用于任何二维单通道图像。 首先，我们可以删除初始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)

Let’s try it out:

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

Out:

0 0.3497849924564362
1 0.23474137868881226

Using your GPU

如果您足够幸运能够访问支持CUDA的GPU（您可以从大多数云提供商那里以每小时0.5美元的价格租用一个GPU），您可以使用它来加速您的代码。 首先检查你的GPU是否在Pytorch中工作：

print(torch.cuda.is_available())

Out:

True

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

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

让我们更新预处理以将批次移动到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)

Out:

0 0.21651759614944457
1 0.1779484712600708

Closing thoughts

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

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

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

torch.nn

• Module：创建一个callable，其行为类似于一个函数，但也可以包含状态（例如神经网络层权重）。 它知道它包含哪些参数，并且可以将所有梯度归零，循环通过它们以进行权重更新等。
• Parameter：张量的包装器，告诉模块它具有在backprop期间需要更新的权重。 仅更新具有requires_grad属性集的张量。
• functional：一个模块（通常按惯例导入到F命名空间中），它包含激活函数，损失函数等，以及层的非有状态版本，如卷积层和线性层。

torch.optim：包含SGD等优化器，可在backward步骤中更新Parameter的权重。

Dataset：具有__len__和__getitem__的对象的抽象接口，包括Pytorch提供的类，如TensorDataset。

DataLoader：获取任何数据集并创建一个返回批量数据的迭代器。