编写高效的PyTorch代码技巧（下）
原文：https://github.com/vahidk/EffectivePyTorch
作者：vahidk
前言
这是一份 PyTorch 教程和最佳实践笔记，目录如下所示：
- PyTorch 基础
- 将模型封装为模块
- 广播机制的优缺点
- 使用好重载的运算符
- 采用 TorchScript 优化运行时间
- 构建高效的自定义数据加载类
- PyTorch 的数值稳定性
上篇文章的链接如下：
这次介绍后面3点，写出高效的代码以及保证做数值计算时候的稳定性。
5. 采用 TorchScript 优化运行时间
PyTorch 优化了维度很大的张量的运算操作。在 PyTorch 中对小张量进行太多的运算操作是非常低效的。所以有可能的话，将计算操作都重写为批次（batch）的形式，可以减少消耗和提高性能。而如果没办法自己手动实现批次的运算操作，那么可以采用 TorchScript 来提升代码的性能。
TorchScript 是一个 Python 函数的子集，但经过了 PyTorch 的验证，PyTorch 可以通过其 just in time(jtt) 编译器来自动优化 TorchScript 代码，提高性能。
下面给出一个具体的例子。在机器学习应用中非常常见的操作就是 batch gather ，也就是 output[i] = input[i, index[i]]。其代码实现如下所示：
import torch
def batch_gather(tensor, indices):
output = []
for i in range(tensor.size(0)):
output += [tensor[i][indices[i]]]
return torch.stack(output)
通过 torch.jit.script 装饰器来使用 TorchScript 的代码：
@torch.jit.script
def batch_gather_jit(tensor, indices):
output = []
for i in range(tensor.size(0)):
output += [tensor[i][indices[i]]]
return torch.stack(output)
这个做法可以提高 10% 的运算速度。
但更好的做法还是手动实现批次的运算操作，下面是一个向量化实现的代码例子，提高了 100 倍的速度：
def batch_gather_vec(tensor, indices):
shape = list(tensor.shape)
flat_first = torch.reshape(
tensor, [shape[0] * shape[1]] + shape[2:])
offset = torch.reshape(
torch.arange(shape[0]).cuda() * shape[1],
[shape[0]] + [1] * (len(indices.shape) - 1))
output = flat_first[indices + offset]
return output
6. 构建高效的自定义数据加载类
上一节介绍了如何写出更加高效的 PyTorch 的代码，但为了让你的代码运行更快，将数据更加高效加载到内存中也是非常重要的。幸运的是 PyTorch 提供了一个很容易加载数据的工具，即 DataLoader 。一个 DataLoader 会采用多个 workers 来同时将数据从 Dataset 类中加载，并且可以选择使用 Sampler 类来对采样数据和组成 batch 形式的数据。
如果你可以随时访问你的数据，那么使用 DataLoader 会非常简单：只需要继承 Dataset 类别并实现 __getitem__ (读取每个数据)和 __len__（返回数据集的样本数量）这两个方法。下面给出一个代码例子，如何从给定的文件夹中加载图片数据：
import glob
import os
import random
import cv2
import torch
class ImageDirectoryDataset(torch.utils.data.Dataset):
def __init__(path, pattern):
self.paths = list(glob.glob(os.path.join(path, pattern)))
def __len__(self):
return len(self.paths)
def __item__(self):
path = random.choice(paths)
return cv2.imread(path, 1)
比如想将文件夹内所有的 jpeg 图片都加载，代码实现如下所示：
dataloader = torch.utils.data.DataLoader(ImageDirectoryDataset("/data/imagenet/*.jpg"), num_workers=8)
for data in dataloader:
# do something with data
这里采用了 8 个 workers 来并行的从硬盘中读取数据。这个数量可以根据实际使用机器来进行调试，得到一个最佳的数量。
当你的数据都很大或者你的硬盘读写速度很快，采用DataLoader进行随机读取数据是可行的。但也可能存在一种情况，就是使用的是一个很慢的连接速度的网络文件系统，请求单个文件的速度都非常的慢，而这可能就是整个训练过程中的瓶颈。
一个更好的做法就是将数据保存为一个可以连续读取的连续文件格式。例如，当你有非常大量的图片数据，可以采用 tar 命令将其压缩为一个文件，然后用 python 来从这个压缩文件中连续的读取图片。要实现这个操作，需要用到 PyTorch 的 IterableDataset。创建一个 IterableDataset 类，只需要实现 __iter__ 方法即可。
下面给出代码实现的例子：
import tarfile
import torch
def tar_image_iterator(path):
tar = tarfile.open(self.path, "r")
for tar_info in tar:
file = tar.extractfile(tar_info)
content = file.read()
yield cv2.imdecode(content, 1)
file.close()
tar.members = []
tar.close()
class TarImageDataset(torch.utils.data.IterableDataset):
def __init__(self, path):
super().__init__()
self.path = path
def __iter__(self):
yield from tar_image_iterator(self.path)
不过这个方法有一个问题，当使用 DataLoader 以及多个 workers 读取这个数据集的时候，会得到很多重复的数据：
dataloader = torch.utils.data.DataLoader(TarImageDataset("/data/imagenet.tar"), num_workers=8)
for data in dataloader:
# data contains duplicated items
这个问题主要是因为每个 worker 都会创建一个单独的数据集的实例，并且都是从数据集的起始位置开始读取数据。一种避免这个问题的办法就是不是压缩为一个tar 文件，而是将数据划分成 num_workers 个单独的 tar 文件，然后每个 worker 分别加载一个，代码实现如下所示：
class TarImageDataset(torch.utils.data.IterableDataset):
def __init__(self, paths):
super().__init__()
self.paths = paths
def __iter__(self):
worker_info = torch.utils.data.get_worker_info()
# For simplicity we assume num_workers is equal to number of tar files
if worker_info is None or worker_info.num_workers != len(self.paths):
raise ValueError("Number of workers doesn't match number of files.")
yield from tar_image_iterator(self.paths[worker_info.worker_id])
所以使用例子如下所示：
dataloader = torch.utils.data.DataLoader(
TarImageDataset(["/data/imagenet_part1.tar", "/data/imagenet_part2.tar"]), num_workers=2)
for data in dataloader:
# do something with data
这是一种简单的避免重复数据的问题。而 tfrecord 则用了比较复杂的办法来共享数据，具体可以查看：
https://github.com/vahidk/tfrecord
7. PyTorch 的数值稳定性
当使用任意一个数值计算库，比如 NumPy 或者 PyTorch ，都需要知道一点，编写数学上正确的代码不一定会得到正确的结果，你需要确保这个计算是稳定的。
首先以一个简单的例子开始。从数学上来说，对任意的非零 x ，都可以知道式子 是成立的。但看看具体实现的时候，是不是总是正确的：
import numpy as np
x = np.float32(1)
y = np.float32(1e-50) # y would be stored as zero
z = x * y / y
print(z) # prints nan
代码的运行结果是打印 nan ，原因是 y 的数值对于 float32 类型来说非常的小，这导致它的实际数值是 0 而不是 1e-50。
另一种极端情况就是 y 非常的大：
y = np.float32(1e39) # y would be stored as inf
z = x * y / y
print(z) # prints nan
输出结果依然是 nan ，因为 y 太大而被存储为 inf 的情况，对于 float32 类型来说，其范围是 1.4013e-45 ~ 3.40282e+38，当超过这个范围，就会被置为 0 或者 inf。
下面是如何查看一种数据类型的数值范围：
print(np.nextafter(np.float32(0), np.float32(1))) # prints 1.4013e-45
print(np.finfo(np.float32).max) # print 3.40282e+38
为了让计算变得稳定，需要避免过大或者过小的数值。这看起来很容易，但这类问题是很难进行调试，特别是在 PyTorch 中进行梯度下降的时候。这不仅因为需要确保在前向传播过程中的所有数值都在使用的数据类型的取值范围内，还要保证在反向传播中也做到这一点。
下面给出一个代码例子，计算一个输出向量的 softmax，一种不好的代码实现如下所示：
import torch
def unstable_softmax(logits):
exp = torch.exp(logits)
return exp / torch.sum(exp)
print(unstable_softmax(torch.tensor([1000., 0.])).numpy()) # prints [ nan, 0.]
这里计算 logits 的指数数值可能会得到超出 float32 类型的取值范围，即过大或过小的数值，这里最大的 logits 数值是 ln(3.40282e+38) = 88.7，超过这个数值都会导致 nan 。
那么应该如何避免这种情况，做法很简单。因为有 ，也就是我们可以对 logits 减去一个常量，但结果保持不变，所以我们选择logits 的最大值作为这个常数，这种做法，指数函数的取值范围就会限制为 [-inf, 0] ，然后最终的结果就是 [0.0, 1.0] 的范围，代码实现如下所示：
import torch
def softmax(logits):
exp = torch.exp(logits - torch.reduce_max(logits))
return exp / torch.sum(exp)
print(softmax(torch.tensor([1000., 0.])).numpy()) # prints [ 1., 0.]
接下来是一个更复杂点的例子。
假设现在有一个分类问题。我们采用 softmax 函数对输出值 logits 计算概率。接着定义采用预测值和标签的交叉熵作为损失函数。对于一个类别分布的交叉熵可以简单定义为 ：
所以有一个不好的实现交叉熵的代码实现为：
def unstable_softmax_cross_entropy(labels, logits):
logits = torch.log(softmax(logits))
return -torch.sum(labels * logits)
labels = torch.tensor([0.5, 0.5])
logits = torch.tensor([1000., 0.])
xe = unstable_softmax_cross_entropy(labels, logits)
print(xe.numpy()) # prints inf
在上述代码实现中，当 softmax 结果趋向于 0，其 log 输出会趋向于无穷，这就导致计算结果的不稳定性。所以可以对其进行重写，将 softmax 维度拓展并做一些归一化的操作：
def softmax_cross_entropy(labels, logits, dim=-1):
scaled_logits = logits - torch.max(logits)
normalized_logits = scaled_logits - torch.logsumexp(scaled_logits, dim)
return -torch.sum(labels * normalized_logits)
labels = torch.tensor([0.5, 0.5])
logits = torch.tensor([1000., 0.])
xe = softmax_cross_entropy(labels, logits)
print(xe.numpy()) # prints 500.0
可以验证计算的梯度也是正确的：
logits.requires_grad_(True)
xe = softmax_cross_entropy(labels, logits)
g = torch.autograd.grad(xe, logits)[0]
print(g.numpy()) # prints [0.5, -0.5]
这里需要再次提醒，进行梯度下降操作的时候需要额外的小心谨慎，需要确保每个网络层的函数和梯度的范围都在合法的范围内，指数函数和对数函数在不正确使用的时候都可能导致很大的问题，它们都能将非常小的数值转换为非常大的数值，或者从很大变为很小的数值。