9.torch.nn是啥 (What is torch.nn really)

[TOC]

PyTorch 提供许多优雅的模块,如 torch.nntorch.optimDatasetDataLoader等,来协助创建和训练网络。为了充分利用这些来解决问题,需要真正地理解它们到底在做些啥。首先,先不用这些模块的任何功能,在 MNIST 数据上训练一个基础的神经网络,一开始仅用 PyTorch 的最基本的 Tensor 函数。然后,逐步添加这些模块的功能,精确地展示每一块完成什么,并如何工作使得代码更加简介,灵活。

MNIST data setup

使用 pathlib 处理路径,使用 requests 下载数据。仅在使用模块时才导入它们,这样就可以明确知道每个点用了什么。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
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)

pathlib 提供表示文件系统路径的类,其语义适合各种操作系统。分成两类,一类是 pure paths,纯粹的计算操作没有 I/O,另一类是 concrete paths,继承自 pure paths 但提供 I/O 操作。
该数据集是 NumPy 数组格式,使用一种特殊的 Python 序列化数据格式 pickle 存储的。使用 gzip 来打开。

1
2
3
4
5
6
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。可视化需要将其 reshape 为2维的。

1
2
3
4
5
from matplotlib import pyplot

print(x_train.shape)
pyplot.imshow(x_train[0].reshape((28, 28)), cmap='gray')
pyplot.show()

PyTorch 使用 torch.tensor 而不是 NumPy 数组,因此需要转换数据。

1
2
3
4
5
6
7
8
9
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())

Neural net from scratch (no torch.nn)

首先使用 PyTorch Tensor 运算来创建一个模型。PyTorch 提供方法创建随机或零填充的 Tensors,将用这些来为线性模型创建权重和偏置。这些就是普通的 Tensor,仅有一项特殊的附加项:告知 PyTorch 它们需要梯度。这使得 PyTorch 记录 Tensor 上的运算,这样就能在反向传播期间自动计算梯度。

在权重初始化后,通过带下划线的 requires_grad_() 设置 requires_grad。(下划线在 PyTorch 中标志运算就地执行),使用 Xavier 初始化权重,乘上 1/sqrt(n)。

1
2
3
4
5
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 代码。

1
2
3
4
5
6
7
def (x):
# 对最后一个维度求和,unsqueeze 表示在指定位置插入一个维度
return x - x.exp().sum(-1).log().unsqueeze(-1)


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

上述代码中,@ 表示点乘运算。在一批数据上调用这些函数(该情况为64张图像)。这是前向传播的一部分。

1
2
3
4
batch_size = 64
xb = x_train[0:batch_size]
preds = model(xb)
print(preds[0], preds.shape)

preds 中不仅包含了 Tensors 数值,还包含了梯度函数,稍后将使用这些函数进行反向传播。实现负对数似然函数来作为损失函数(仍然仅用标准的 Python),然后用上面的随机模型来查看一下损失,稍后就可以看到在反向传播后是否有所提升。

1
2
3
4
5
6
7
def nll(inputs, target):
return -inputs[range(target.shape[0]), target].mean()


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

同时也实现一个函数来计算模型的准确率,如果最大值的索引匹配目标值,则预测是正确的。并查看一下上面的随机模型的准确率。

1
2
3
4
5
6
def accuracy(out, yb):
preds = torch.argmax(out, dim=1)
return (preds == yb).float().mean()


print(accuracy(preds, yb))

现在可以运行训练循环。每次迭代都将:

  • 选择小批次数据(大小为 batch_size )
  • 使用模型作预测
  • 计算损失
  • loss.backward() 更新模型的梯度

现在使用这些梯度来更新 weightsbias。在 torch.no_grad() 上下文管理器中做这些,因为不需要为下次梯度计算来记录这次的更新参数行为。之后将这些梯度置零,准备下一循环。否则,梯度将记录所发生的所有运算(无论已经存储了什么,loss.backward() 都会叠加梯度,而不是覆盖它们)

可以使用标准的 Python 调试器来单步调试 PyTorch 代码,在每一步检查变量值。取消注释 set_trace() 来尝试。迭代后损失应该降低,准确率上升。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from IPython.core.debugger import set_trace

lr = 0.5 # 学习率
epochs = 2 # 训练多少轮

for epoch in range(epochs):
for i in range((n - 1) // batch_size + 1):
# set_trace()
start_i = i * batch_size
end_i = start_i + batch_size
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))

Using torch.nn.functional

使用 PyTorch 的 nn 类重构之前的代码,使得代码更加简介和灵活。第一步也是最简单的一步就是使用 torch.nn.functional (通常按照约定导入为命名空间 F )的函数重构之前手撸的激活函数和损失函数。这个模块包含 torch.nn 库中所有的函数(该库剩余的部分包含类)。该库不仅有许多损失和激活函数,还有一些实用的函数,例如池化函数。(还有些函数用来做卷积,线性层等等,但这些通常用 torch.nn 库来做更好)

如果使用负对数似然损失函数和 log softmax 激活函数,PyTorch 提供一个单独的函数 F.cross_entropy 来组合它们。下面使用它们来定义模型,并查看一下损失和准确率。

1
2
3
4
5
6
7
8
import torch.nn.functional as F

loss_func = F.cross_entropy

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

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

Refactor using nn.Module

使用 nn.Modulenn.Parameter 以获得更清晰和更简洁的训练循环。继承 nn.Module 实现子类(能够跟踪状态)。nn.Module 有许多的属性和方法(例如 .parameters().zero_grad() )可以使用。

1
2
3
4
5
6
7
8
9
10
11
12
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 方法。还是跟以前一样计算损失。

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

现在可以使用 model.parameters()model.zero_grad() (都是由 PyTorch 为 nn.Module 定义的)使得更新参数更加简洁,并且更不容易忘记一些参数,尤其是对于复杂模型:

1
2
3
4
with torch.no_grad():
for p in model.parameters():
p -= p.grad * lr
model.zero_grad()

将训练循环包装在 fit 函数内:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def fit():
for epoch in range(epochs):
for i in range((n - 1) // batch_size - 1):
start_i = i * batch_size
end_i = start_i + batch_size
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))

Refactor using nn.Linear

这里使用 PyTorch 的 nn.Linear 类来实现线性层。PyTorch 有许多预先定义好的层,能极大简化代码,并且通常更快。

1
2
3
4
5
6
7
8
9
10
11
12
class MnistLogistic(nn.Module):
def __init__(self):
super().__init__()
self.lin = nn.Linear(784, 10)

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

model = MnistLogistic()
print(loss_func(model(xb), yb))
fit()
print(loss_func(model(xb), yb))

Refactor using optim

PyTorch 还有个带有各种优化算法的包 torch.optim。可以使用优化器的方法 .step() 来更新全部参数,而不用手动更新每个参数。

1
2
optimizer.step()
optimizer.zero_grad()

注意需要在下次计算梯度之前将梯度缓存清空,.zero_grad() 方法清空梯度的缓存。下面定义一个函数来创建模型和优化器:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from torch import optim

def get_model():
model = MnistLogistic()
return model, optim.SGD(model.parameters(), lr=lr)

model, optimizer = get_model()
print(loss_func(model(xb), yb))

for epoch in range(epochs):
for i in range((n - 1) // batch_size - 1):
start_i = i * batch_size
end_i = start_i + batch_size
xb = x_train[start_i:end_i]
yb = y_train[start_i:end_i]
pred = model(xb)
loss = loss_func(pred, yb)
loss.backward()

optimizer.step()
optimizer.zero_grad()

Refactor using Dataset

PyTorch 有一个抽象的 Dataset 类,具有 __len__ 函数(被 Python 的 len 函数调用)和作为其索引方式的 __getitem__ 函数。

PyTorch 的 TensorDataset 是一个封装 Tensor 的 Dataset。通过定义长度和索引的方式,可以迭代,索引和沿着 tensor 的第一个维度切片,使得在训练过程中访问自变量和因变量更加容易。

1
2
3
4
from torch.utils.data import TensorDataset

train_ds = TensorDataset(x_train, y_train)
xb, yb = train_ds[i*batch_size:(i+1)*batch_size]

x_trainy_train 组合在单个 TensorDataset 中,更加容易迭代和切片。

Refactor using DataLoader

PyTorch 的 DataLoader 负责管理批处理,可以从任意的 Dataset 创建 DataLoaderDataLoader 使得迭代批处理更加容易。DataLoader 自动给出每个小批量处理的数据,而不是使用切片。

1
2
3
4
5
6
from torch.utils.data import DataLoader

train_dl = DataLoader(train_ds, batch_size=batch_size)
for xb, yb in train_dl:
pred = model(xb)
loss = loss_func(pred, yb)

由于 PyTorch 的 nn.Modulenn.ParameterDatasetDataLoader 等,训练循环已经显著精简了许多,并且更容易理解。接下来尝试添加实际中创建高效模型所必需的基本特性。

Add validation

实际上,还应该有验证集来检验是否过拟合。打乱训练数据对预防批量数据之间的相关性和过拟合很重要,另外,无论是否打乱验证集,验证集上的损失都应该相同,并且由于打乱数据需要额外的时间,因此打乱验证集没有意义。

将验证集的批量大小设为训练集的两倍。这是因为验证集不需要反向传播,因此占用的内存更少(不需要存储梯度),所以可以设更大的批量大小,并且更快地计算损失。

1
2
3
4
5
train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)

valid_ds = TensorDataset(x_valid, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=2 * batch_size)

在每轮结束计算和打印验证集的损失。在训练之前先调用 model.train(),而在推理前先调用 model.eval(),这是因为这些会被如 nn.BatchNorm2dnn.Dropout 所使用来确保它们在不同阶段的行为。(训练阶段需要,而评估阶段不需要)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
model, optimizer = 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()
optimizer.step()
optimizer.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))

Create fit() and get_data()

将计算训练集损失和验证集损失合在一起,编写一个函数计算两者,loss_batch,该函数计算批处理的损失。在计算训练集时,传递优化器来进行反向传播,而验证集不需要。

1
2
3
4
5
6
7
8
9
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 运行必须的运算来训练模型,并且计算每轮训练和验证损失。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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]
)
valid_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

print(epoch, valid_loss)

get_data 返回关于训练集和验证集的 DataLoaders

1
2
3
4
5
def get_data(train_ds, valid_ds, batch_size):
return (
DataLoader(train_ds, batch_size=batch_size, shuffle=True),
DataLoader(valid_ds, batch_size=2 * batch_size),
)

现在,整个获取数据和拟合模型的过程在3行代码完成:

1
2
3
train_dl, valid_dl = get_data(train_ds, valid_ds, batch_size)
model, optimizer = get_model()
fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)

这三行代码可训练各种模型,下面使用它们来训练一个卷积神经网络。

Switch to CNN

现在准备构建一个三层卷积神经网络,因为前面部分的函数没有假定模型的形式,所以可以用来训练 CNN 而不用做任何修改。使用 PyTorch 预先定义的 Conv2d 作为卷积层,来定义一个有3层卷积层的 CNN。每个卷积层都紧跟着 ReLu,最后再进行平均池化。(注意 view 是 NumPy 的 reshape 的 PyTorch 版本)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class MnistCNN(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 = MnistCNN()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)

nn.Sequential

torch.nn 还有另外一个方便的类可以用来简化代码:SequentialSequential 对象按顺序地运行每个包含在它内部的模块。需要能够从给定的函数轻易地定义些自定义的层,如 PyTorch 不提供的 view 层,Lambda 将创建一个层,在使用 Sequential 创建网络时使用。

1
2
3
4
5
6
7
8
9
10
11
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 创建模型非常简单:

1
2
3
4
5
6
7
8
9
10
11
12
13
model = nn.Sequential(
Lambda(preprocess),
nn.Conv2d(1, 16, 3, 2, 1),
nn.ReLU(),
nn.Conv2d(16, 16, 3, 2, 1),
nn.ReLU(),
nn.Conv2d(16, 10, 3, 2, 1),
nn.ReLU(),
nn.AvgPool2d(4),
Lambda(lambda x: x.view(x.size(0), -1)),
)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)

Wrapping DataLoader

定义的 CNN 非常简洁,但只适用于 MNIST,因为:

  • 假设输入是28*28的长向量
  • 假设最后的 CNN 网格大小为4*4

摆脱这两个假设,使模型可以在任意2维的单通道图像上运行。首先从初始化移除 Lambda 层,将其移入到生成器。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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,这允许定义想要的输出大小,而不是取决于输入。因此,模型可以在任意大小的输入上运行。

1
2
3
4
5
6
7
8
9
10
11
12
13
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)),
)

optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)

Using your GPU

首先检查是否有可用的 GPU:

1
print(torch.cuda.is_availabel())

然后创建一个 device 对象:

1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

最后将模型,预处理等移至 GPU 上,然后进行拟合

1
2
3
4
5
def preprocess(x, y):
return x.view(-1, 1, 28, 28).to(device), y.to(device)

model.to(device)
fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)

Closing thoughts

现在有了能够使用 PyTorch 训练各种模型的通用的数据管道和训练循环。然而,还有许多东西哇,如数据增强,超参数调整,监视训练过程,迁移学习等等。这些功能在 fastai 库可用,使用同样的设计方法开发。

总结:

  • torch.nn
    • Module:创建可像函数一样调用的,但同样包含状态的(例如神经网络层权重)模块。它知道自己包含的参数,并且能够清零梯度,循环遍历以更新权重等
    • Parameter:对 Tensor 的封装,它告知 Module 有权重需要在反向传播中更新,仅 requires_grad 为True 的 Tensor 才需要更新
    • functional:一个包含激活函数,损失函数等的模块(通常被作为命名空间 F 导入),还是如卷积层和线性层等的无状态版本
  • torch.optim:包含如 SGD 等优化器,在反向传播期间更新 Parameter 的权重
  • Dataset:一个有 __len____getitem__ 的对象的抽象接口,包括 PyTorch 提供的类如 TensorDataset
  • DataLoader:接受任意 Dataset,并且创建一个返回批量数据的迭代器