torch.nn 的本质

torch.nn 的本质

PyTorch 提供了各种优雅设计的 modules 和类 torch.nn，torch.optim，Dataset 和 DataLoader 来帮助你创建并训练神经网络。为了充分利用它们的力量并且根据你的问题定制它们，你需要真正地准确了解它们在做什么。为了建立这种理解，我们首先从这些模型（models）上不使用任何特性（features）在 MNIST 数据集上训练一个基本的神经网络；我们将从最基本的 PyTorch Tensor 功能开始。然后，我们每次在 torch.nn，torch.optim，Dataset 或 DataLoader 逐渐地增加一个特性，准确地展示每一块做的事情，并且它如何使代码更简洁或更灵活。

这篇博文假设你已经安装了 PyTorch 并且熟悉 Tensor 操作的基础。（如果你熟悉 NumPy 数组的操作，你会发现这里使用的 PyTorch Tensor 操作几乎相同。）

MNISt 数据配置

我们将使用经典的 MNIST 数据集，其是由手写数字（从 0 到 9）的黑白图像组成。

我们将使用 pathlib 处理路径（Python3 的标准库之一），使用 request 下载数据集。我们在每一步仅导入使用的 modules，所以你可以准确地看到每一步在使用什么。

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)

如果网速不给力，可以从这里下载我下载好的 mnist.pkl.gz。

这个数据集是存储在 NumPy 数组的格式，而且已经被 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')

每一张图像是 28x28，并且被存储为展开的长度为 784（=28x28）的一行。让我们看一个，首先，我们需要重新将形状（shape）改为二维的。

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 数组，所以我们需要转换我们的数据。

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 Tensot 之外的包开始构建一个模型。我们假设你已经熟悉神经网络的基础。（如果你还不熟悉，你可从 course.fast.ai 学习它们）。

PyTorch 提供了创建随机数或零值填充 Tensor 的方法，我们将使用它创建我们简单线性模型的权重（weight）和偏置单元（bias）。这些只是普通的 Tensor，但是一个非常特殊的附加：我们告诉 PyTorch 它们需要梯度。这让 PyTorch 记录所有完成在 Tensor 上的操作，以便它在反向传播时自动地计算梯度！

对于这些权重（weight），我们初始化之 后 设置 requires_grad，因为我们不希望这个步骤（初始化）被添加进梯度。（注意下划线符号 _，在 PyTorch 中表明某个 Tensor 上的操作就地执行（in-place）。）

这里我们使用 Xavier initialisation（通过乘以 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 object））作为模型！所以，让我们仅仅使用简单的矩阵相乘和广播（broadcasted）加法创建一个线性模型。我们也需要一个激活函数（activation function），所以我们将写一个 log_softmax 使用。记住：即使 PyTorch 提供了许多写好的损失函数（loss function）、激活函数（activation function）等等，你也可以使用原生的 Python 写出你自己的函数。甚至 PyTorch 还将为你的函数自动地创建快速的 GPU 或矢量化（vectorized）CPU 代码。

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)

在上面的代码中，@ 符号表示点积（dot product）操作。我们将在一个数据批量上调用我们的函数（在这个例子中，64 张图片），这是一次前向传播（forward pass）。注意在这个阶段我们的预测不比随即预测好，因为我们是从随机权重开始的。

bs = 64  # batch size

xb = x_train[0:bs]  # a mini-batch from x
preds = model(xb)
print(preds[0], preds.shape)

tensor([-2.6015, -2.8883, -3.1596, -2.2470, -2.8118, -2.0224, -2.2773, -2.1566,
        -1.4275, -2.6397], grad_fn=<SelectBackward>) torch.Size([64, 10])

正如你所见，preds Tensor 包含了不仅仅是 Tensor 中的值，同样也有一个梯度函数。我们之后将使用它做反向传播。

让我们实现负对数似然（negative log-likelihood）作为我们的损失函数（再次说明，我们只使用原生的 Python）。

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

loss_func = nll

让我们查看我们的随机模型的损失值（loss），之后我们经过反向传播之后看看是否得到了提升。

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

tensor(2.4096, grad_fn=<NegBackward>)

让我们实现一个函数计算我们模型的准确率。对于每一个预测，如果最大值的下标（index）和目标值一样，那么预测就是正确的。

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

同样检查我们随机模型的准确率，并且在反向传播之后查看准确率是否得到了提升。

print(accuracy(preds, yb))

tensor(0.0625)

我们现在可以执行训练循环。对于每一迭代（iteration），我们将：

• 选择一个数据的批量（大小为 bs）
• 使用模型做预测
• 计算损失值（loss）
• loss.backward() 更新模型的梯度，在这个例子中，是 weight 和 bias。

我们使用这些梯度更新权重（weight）和偏移（bias）。我们在 torch.no_grad() 上下文管理器内做更新，因为我们不希望这些活动被记录在我们的下一步梯度的计算。你可以在 这里 查看更多关于 PyTorch 的 autograd 记录操作。

下一步我们将梯度设为 0，以便为下一个循环准备。否则，我们的梯度会记录所有已经发生的运算的运行记录（比如 loss.backward() 会累加梯度，无论里面存储了什么，而不是替换）。

你可以使用标准的 Python 调试器（debugger）单步调试（step through）PyTorch 的代码，让你可以检查每一个步骤的变量值。

取消注释 set_trace() 尝试它。

from IPython.core.debugger import set_trace

lr = 0.5  # 学习率（learning rate）
epochs = 2  # 训练多少次

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

我们已经完全地从零开始创建并训练了一个最小的神经网络（在这个例子中，一个逻辑回归（logistic regression），因为我们没有隐藏层）。

让我们检查损失值（loss）和准确率与我们之前得到的相比。我们期待损失值（loss）将会下降并且准确率将会有所上升。

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

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

使用 torch.nn.functional

我们下一步就重构（refactor）我们的代码，以便它和之前做的一样，只有我们开始利用 PyTorch 的 nn 类使代码变得更加简洁和灵活。从这里开始的每一步，我们应该使我们的代码变得一个或多个的：简短、更容易理解或更灵活。

在一开始，最简单的步骤是通过替换我们手写的激活函数（activate function）和损失函数（loss function）为 torch.nn.functional 包中的函数（依照惯例，通常我们导入到命名空间（namespace）F 中），让我们的代码变得更简短。这个 module 包含了 torch.nn 库内的所有函数（而该库的其它部分包含了类（classes））。以及各种各样的损失（loss）和激活（activation）函数，你也在这里可以找到一些方便的函数来构建神经网络，比如池化函数（pooling functions）。（也有做卷积（convolutions）的函数、线性层（linear layers）等等，但是我们即将看到，这些通常使用库的其它部分更好地处理。）

如果你正使用负对数似然（negative log likelihood）损失函数和 log softmax 函数，PyTorch 提供了单一的函数 F.cross_entropy 将二者结合起来。所以我们甚至可以从我们的模型移除激活函数（activation function）。

import torch.nn.functional as F

loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

注意我们不再在 model 函数里调用 log_softmax 函数。让我们确认我们的损失值（loss）和准确率是否和之前一样。

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

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

使用 torch.Module 重构

下一步，为了更清楚和更简洁的训练循环（training loop），我们将使用 nn.Module 和 nn.Parameter。我们的子类 nn.Module（它本身是一个类并且可以跟踪状态）。在这个例子，我们想要创建一个持有权重（weights）、偏移（bias）和前向传播的方法的类。nn.Module 有一些我们将使用的属性和方法（比如 .parameters() 和 .zero_grad()）。

nn.Module（大写 M）是 PyTorch 特有的概念，并且是一个我们将经常使用的类。nn.Module 不要与 Python 的 module（小写 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.randn(10))

    def forward(self, xb):
        return xb @ self.weights + self.bias

因为我们现在使用的是对象而不是使用函数，所以我们需要先实例化我们的模型。

model = Mnist_Logistic()

现在我们可以和以前一样以相同的方法计算损失值（loss）。注意 nn.Module 的对象好像和函数一样使用（即它们是可调用的（callable），而且在后台 PyTorch 将自动调用我们的方法 forward。

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

tensor(3.0925, grad_fn=<NllLossBackward>)

之前对于我们的训练循环来说，我们必须通过变量名来更新每一个参数的值，并且要单独地对每一个参数的梯度手动清零，就像这样。

with torch.no_grad():
    weights -= weights.grad * lr
    bias -= bias.grad * lr
    weights.grad_zeor_()
    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()

让我们再次检查我们的损失值（loss）有所减少。

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

tensor(0.0814, grad_fn=<NllLossBackward>)

使用 torch.Linear 重构

我们继续重构我们的代码。作为手动定义和初始化 self.weights 和 self.bias 并且计算 xb @ self.weights + self.bias 的替代，我们将用 PyTorch 的类 nn.Linear 为一个为我们做所有事情的线性层（linear layer）。PyTorch 有许多层（layers）的类型，可以极大地简化我们的代码，同样也使其更快。

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

    def forward(self, xb):
        return self.lin(xb)

我们实例化我们的模型并和以前同样的方法计算损失值（loss）。

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

tensor(2.3702, grad_fn=<NllLossBackward>)

我们仍然能够使用之前的 fit 方法。

fit()

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

tensor(0.0813, grad_fn=<NllLossBackward>)

使用 optim 重构

PyTorch 同样有包含各种优化算法的包 torch.optim。我们可以从我们的优化器（optimizer）使用 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

我们将定义一个小函数来创建我们的模型和优化器（optimizer）让我们在之后可以重复使用它。

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.2597, grad_fn=<NllLossBackward>)
tensor(0.0809, grad_fn=<NllLossBackward>)

使用 Dataset 重构

PyTorch 有一个抽象 Dataset 类。Dataset 可以具有 __len__ 函数（通过 Python 的标准 len 函数调用）和 __getitem__ 函数作为对其索引的一种方法。这个例子 是一个非常好的例子来创建一个定制的继承 Dataset 的 FacialLandmarkDataset 类。

PyTorch 的 TensorDataset 是一个封装了 Dataset 的 Tensors。通过定义索引的长度和方式，这同样给我们沿着 Tensor 的第一个维度迭代、索引和切片（slice）的方法。这让我们训练时更容易在同一行中访问自变量和因变量。

from torch.utils.data import TensorDataset

x_train 和 y_train 都可以被绑定在一个更容易迭代和切片（slice）的 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.0823, grad_fn=<NllLossBackward>)

使用 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）从 DataLoader 自动地加载：

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.0825, grad_fn=<NllLossBackward>)

感谢 PyTorch 的 nn.Module，nn.Parameter，Dataset 和 DataLoader，我们的训练循环现在显著的小并且非常容易理解。让我们现在尝试增加在实际中创建高效的模型的基本特征。

增加验证（Add validation）

在第一部分，我们只是尝试建立一个合理的训练循环以用于我们的训练数据。在实际上，你总是应该有一个 验证集（validation set），为了鉴别你是否过拟合（overfitting）。

洗乱（shuffling)训练数据对于防止数据批量和过拟合之间的相关性（correlation）很 重要。在另一方面，无论我们洗乱（shuffle）验证集与否，验证损失（validation loss）都是一样的。由于洗乱（shuffling）花费额外的时间，洗乱（shuffle）验证数据是没有意义的。

我们将设置验证集（validation set）的批量大小为训练集的两倍。这是因为验证集不需要反向传播并且占用更少的内存（它不需要存储梯度）。我们利用这一点使用大的数据批量并且更快地计算损失值（loss）。

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() 而且在评估（inference）之前调用 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.3125)
1 tensor(0.2864)

创建 fit() 和 get_data()

我们现在对我们自己进行一些小的重构。因为我们经历了两次相似的计算训练集（training set）和验证集（validation set）的损失值（loss）的处理过程，让我们把它变成它自己的函数，loss_batch 来计算一个数据批量的损失值（loss）。

当是训练集时，我们传入一个优化器（optimizer）并且用它做反向传播。对于验证集（validation set），我们不需要传入优化器（optimizer），所以方法（method）不需要执行反向传播。

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）计算训练（training）和验证（validation）的损失（loss）。

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 返回训练集和验证集的 DataLoader。

def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs*2)
    )

现在，我们的整个获取 DataLoader 和训练模型的过程可以运行在 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.3818369417190552
1 0.29548657131195066

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

切换到 CNN

我们现在要构建一个三层卷积层（convolutional layer）的神经网络。因为前面部分的没有一个函数显露出有关模型形式的信息，所以我们将可以使用它们不做任何修改训练一个卷积神经网络（Convolutional Neural Network（CNN））。

我们将使用 PyTorch 的预定义的（predefined）Conv2d 类作为我们的卷积层。我们定义一个有 3 层卷积层的卷积神经网络。每一个卷积后都跟一个 ReLU。在最后，我们执行一个平均池化（average pooling）。（注意 view 是 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 是一种随机梯度下降（stochastic gradient descent）的变体，把之前的更新也考虑在内并且通常让训练更快。

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.38749439089894294
1 0.2610516972362995

nn.Sequential

torch.nn 有另一个方便的类我们可以使用简化我们的代码：Sequential。一个 Sequential 对象以一种顺序的方式运行包含在它之内的 Modules。这是一种写神经网络更简单的方式。

为了利用它，我们需要可以从一个给定的函数简单地定义一个 定制层（custom layer）。举个例子，PyTorch 没有 view 层，我们需要为我们的神经网络创建一个。Lambda 将创建一层（layer），我们可以在使用 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.3955023421525955
1 0.23224713450670242

封装 DataLoader

我们的卷积神经网络相当简洁，但是它只能运行在 MNIST 上，因为：

• 它假设输入是一个 28*28 的长向量
• 它假设最终的卷积神经网络的网格尺寸是 44（因为我们使用平均池化（average pooling）的内核尺寸（kernel size*））

让我们丢掉这两个假设，所以我们的模型可以运行在任何的二维单通道图像上。首先，我们可以移除最开始的 Lambda 层，但是移动数据预处理到一个生成器（generator）。

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.AvgPool2d 为 nn.AdaptiveAvgPool2d，这允许我们定义我们想要的 Tensor 的输出尺寸，而不是我们有的输入 Tensor。因此，我们的模型可以运行在任何尺寸的输入上。

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, 16, 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.4092831357955933
1 0.3001906236886978

使用你的 GPU

如果你足够幸运可以使用一个支持 CUDA（CUDA-capable）的 GPU（你可以以一小时 0.5 刀的价格从很多云提供商租一个）你可以使用它加速你的代码。首先在 PyTorch 里检查你的 GPU 是否可以工作。

print(torch.cuda.is_available())

False

然后为它创建一个设备对象（device object）。

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.18722925274372101
1 0.21267506906986236

总结

我们现在有一套通用的数据通道和训练循环，你可以使用 PyTorch 训练许多模型的类型。

当然，有很多你想要去添加的事情，比如数据增强（data augmentation）、超参调节（hyperparameter tuning）和转移学习（transfer learning）等等。这些特征在 fastai 库都是可用的，这个库已被开发为和这篇博文展示的相同的设计方法，为从业人员进一步提升他们的模型提供了自然的下一步。

我们在这篇博文开始的时候保证过我们通过每一个例子解释 torch.nn，torch.optim，Dataset 和 DataLoader。所以让我们总结一下我们已经看到的。

• torch.nn
  • Module：创建可调用对象（callable）其行为就像一个函数，但是也可以包含状态（state）（比如神经网络层上的权重（weight））。它知道其包含的 Parameter，并且可以清零所有的梯度，遍历它们进行权重更新等。
  • Parameter：一个 Tensor 的包装，用于告诉 Module 它是权重（weight）需要在反向传播时更新。只有设置了 *requires_grad` 属性的 Tensor 才可以被更新。
  • functional：一个模块（module），通常导入转换到 F 的命名空间，它包含激活函数（activation function）、损失函数（loss function）等，以及诸如卷积层和线性层之类的无状态版本。
• torch.optim：包含诸如 SGD 的优化器（optimizer），在反向传播的时候更新 Parameter 的权重（weight）。
• Dataset：一个带有 __len__ 和 __getitem__ 的对象的抽象接口（abstract interface），包含 PyTorch 提供的类，例如 TensorDataset。
• DataLoader：接受任何的 Dataset 并且创建一个返回一个数据批量的迭代器（iterator）。