高性能PyTorch是如何炼成的？过来人吐血整理的10条避坑指南

作者：Eugene Khvedchenya

机器之心编译

参与：小舟、蛋酱、魔王

高性能 PyTorch 的训练管道是什么样的？是产生最高准确率的模型？是最快的运行速度？是易于理解和扩展？还是容易并行化？答案是，包括以上提到的所有。

如何用最少的精力，完成最高效的 PyTorch 训练？一位有着 PyTorch 两年使用经历的 Medium 博主最近分享了他在这方面的 10 个真诚建议。

在 Efficient PyTorch 这一部分中，作者提供了一些识别和消除 I/O 和 CPU 瓶颈的技巧。第二部分阐述了一些高效张量运算的技巧，第三部分是在高效模型上的 debug 技巧。

在阅读这篇文章之前，你需要对 PyTorch 有一定程度的了解。

好吧，从最明显的一个开始：

建议 0：了解你代码中的瓶颈在哪里

命令行工具比如 nvidia-smi、htop、iotop、nvtop、py-spy、strace 等，应该成为你最好的伙伴。你的训练管道是否受 CPU 约束？IO 约束？GPU 约束？这些工具将帮你找到答案。

这些工具你可能从未听过，即使听过也可能没用过。没关系。如果你不立即使用它们也可以。只需记住，其他人可能正在用它们来训练模型，速度可能会比你快 5%、10%、15%-…… 最终可能会导致面向市场或者工作机会时候的不同结果。

数据预处理

几乎每个训练管道都以 Dataset 类开始。它负责提供数据样本。任何必要的数据转换和扩充都可能在此进行。简而言之，Dataset 能报告其规模大小以及在给定索引时，给出数据样本。

如果你要处理类图像的数据（2D、3D 扫描），那么磁盘 I/O 可能会成为瓶颈。为了获取原始像素数据，你的代码需要从磁盘中读取数据并解码图像到内存。每个任务都是迅速的，但是当你需要尽快处理成百上千或者成千上万个任务时，可能就成了一个挑战。像 NVidia 这样的库会提供一个 GPU 加速的 JPEG 解码。如果你在数据处理管道中遇到了 IO 瓶颈，这种方法绝对值得一试。

还有另外一个选择，SSD 磁盘的访问时间约为 0.08–0.16 毫秒。RAM 的访问时间是纳秒级别的。我们可以直接将数据存入内存。

建议 1：如果可能的话，将数据的全部或部分移至 RAM。

如果你的内存中有足够多的 RAM 来加载和保存你的训练数据，这是从管道中排除最慢的数据检索步骤最简单的方法。

这个建议可能对云实例特别有用，比如亚马逊的 p3.8xlarge。该实例有 EBS 磁盘，它的性能在默认设置下非常受限。但是，该实例配备了惊人的 248Gb 的 RAM。这足够将整个 ImageNet 数据集存入内存了！你可以通过以下方法达到这一目标：

class RAMDataset(Dataset):
  def __init__(image_fnames, targets):
    self.targets = targets
    self.images = []
    for fname in tqdm(image_fnames, desc="Loading files in RAM"):
      with open(fname, "rb") as f:
        self.images.append(f.read())

  def __len__(self):
    return len(self.targets)

  def __getitem__(self, index):
    target = self.targets[index]
    image, retval = cv2.imdecode(self.images[index], cv2.IMREAD_COLOR)
    return image, target

我个人也面对过这个瓶颈问题。我有一台配有 4x1080Ti GPUs 的家用 PC。有一次，我采用了有 4 个 NVidia Tesla V100 的 p3.8xlarge 实例，然后将我的训练代码移到那里。鉴于 V100 比我的 oldie 1080Ti 更新更快的事实，我期待看到训练快 15–30%。出乎意料的是，每个时期的训练时间都增加了。这让我明白要注意基础设施和环境差异，而不仅仅是 CPU 和 GPU 的速度。

根据你的方案，你可以将每个文件的二进制内容保持不变，并在 RAM 中进行即时解码，或者对未压缩的图像进行讲解码，并保留原始像素。但是无论你采用什么方法，这里有第二条建议：

建议 2：解析、度量、比较。每次你在管道中提出任何改变，要深入地评估它全面的影响。

假设你对模型、超参数和数据集等没做任何改动，这条建议只关注训练速度。你可以设置一个魔术命令行参数（魔术开关），在指定该参数时，训练会在一些合理的数据样例上运行。利用这个特点，你可以迅速解析管道。

# Profile CPU bottlenecks
python -m cProfile training_script.py --profiling

# Profile GPU bottlenecks
nvprof --print-gpu-trace python train_mnist.py

# Profile system calls bottlenecks
strace -fcT python training_script.py -e trace=open,close,read

Advice 3: *Preprocess everything offline*

建议 3：离线预处理所有内容

如果你要训练由多张 2048x2048 图像制成的 512x512 尺寸图像，请事先调整。如果你使用灰度图像作为模型的输入，请离线调整颜色。如果你正在进行自然语言处理（NLP），请事先做分词处理（tokenization），并存入磁盘。在训练期间一次次重复相同的操作没有意义。在进行渐进式学习时，你可以以多种分辨率保存训练数据的，这还是比线上调至目标分辨率更快。

对于表格数据，请考虑在创建 Dataset 时将 pd.DataFrame 目标转换为 PyTorch 张量。

建议 4：调整 DataLoader 的工作程序

PyTorch 使用一个 DataLoader 类来简化用于训练模型的批处理过程。为了加快速度，它可以使用 Python 中的多进程并行执行。大多数情况下，它可以直接使用。还有几点需要记住：

每个进程生成一批数据，这些批通过互斥锁同步可用于主进程。如果你有 N 个工作程序，那么你的脚本将需要 N 倍的 RAM 才能在系统内存中存储这些批次的数据。具体需要多少 RAM 呢？

我们来计算一下：

1. 假设我们为 Cityscapes 训练图像分割模型，其批处理大小为 32，RGB 图像大小是 512x512x3（高、宽、通道）。我们在 CPU 端进行图像标准化（稍后我将会解释为什么这一点比较重要）。在这种情况下，我们最终的图像 tensor 将会是 512 * 512 * 3 * sizeof(float32) = 3,145,728 字节。与批处理大小相乘，结果是 100,663,296 字节，大约 100Mb；
2. 除了图像之外，我们还需要提供 ground-truth 掩膜。它们各自的大小为（默认情况下，掩膜的类型是 long，8 个字节）——512 * 512 * 1 * 8 * 32 = 67,108,864 或者大约 67Mb；
3. 因此一批数据所需要的总内存是 167Mb。假设有 8 个工作程序，内存的总需求量将是 167 Mb * 8 = 1,336 Mb。

听起来没有很糟糕，对吗？当你的硬件设置能够容纳提供 8 个以上的工作程序提供的更多批处理时，就会出现问题。或许可以天真地放置 64 个工作程序，但是这将消耗至少近 11Gb 的 RAM。

当你的数据是 3D 立体扫描时，情况会更糟糕。在这种情况下，512x512x512 单通道 volume 就会占 134Mb，批处理大小为 32 时，8 个工作程序将占 4.2Gb，仅仅是在内存中保存中间数据，你就需要 32Gb 的 RAM。

对于这个问题，有个能解决部分问题的方案——你可以尽可能地减少输入数据的通道深度：

1. 将 RGB 图像保持在每个通道深度 8 位。可以轻松地在 GPU 上将图像转换为浮点形式或者标准化。
2. 在数据集中用 uint8 或 uint16 数据类型代替 long。

class MySegmentationDataset(Dataset):
  ...
  def __getitem__(self, index):
    image = cv2.imread(self.images[index])
    target = cv2.imread(self.masks[index])

    # No data normalization and type casting here
    return torch.from_numpy(image).permute(2,0,1).contiguous(),
           torch.from_numpy(target).permute(2,0,1).contiguous()

class Normalize(nn.Module):
    # https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/modules/normalize.py
    def __init__(self, mean, std):
        super().__init__()
        self.register_buffer("mean", torch.tensor(mean).float().reshape(1, len(mean), 1, 1).contiguous())
        self.register_buffer("std", torch.tensor(std).float().reshape(1, len(std), 1, 1).reciprocal().contiguous())

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return (input.to(self.mean.type) - self.mean) * self.std

class MySegmentationModel(nn.Module):
  def __init__(self):
    self.normalize = Normalize([0.221 * 255], [0.242 * 255])
    self.loss = nn.CrossEntropyLoss()

  def forward(self, image, target):
    image = self.normalize(image)
    output = self.backbone(image)

    if target is not None:
      loss = self.loss(output, target.long())
      return loss

    return output

通过这样做，会大大减少 RAM 的需求。对于上面的示例。用于高效存储数据表示的内存使用量将为每批 33Mb，而之前是 167Mb，减少为原来的五分之一。当然，这需要模型中添加额外的步骤来标准化数据或将数据转换为合适的数据类型。但是，张量越小，CPU 到 GPU 的传输就越快。

DataLoader 的工作程序的数量应该谨慎选择。你应该查看你的 CPU 和 IO 系统有多快，你有多少内存，GPU 处理数据有多快。

多 GPU 训练 & 推理

神经网络模型变得越来越大。今天，使用多个 GPU 来增加训练时间已成为一种趋势。幸运的是，它经常会提升模型性能来达到更大的批处理量。PyTorch 仅用几行代码就可以拥有运行多 GPU 的所有功能。但是，乍一看，有些注意事项并不明显。

model = nn.DataParallel(model) # Runs model on all available GPUs

运行多 GPU 最简单的方法就是将模型封装在 nn.DataParallel 类中。除非你要训练图像分割模型（或任何生成大型张量作为输出的其他模型），否则大多数情况下效果不错。在正向推导结束时，nn.DataParallel 将收集主 GPU 上所有的 GPU 输出，来通过输出反向运行，并完成梯度更新。

于是，现在就有两个问题：

• GPU 负载不平衡；
• 在主 GPU 上聚合需要额外的视频内存

首先，只有主 GPU 能进行损耗计算、反向推导和渐变步骤，其他 GPU 则会在 60 摄氏度以下冷却，等待下一组数据。

其次，在主 GPU 上聚合所有输出所需的额外内存通常会促使你减小批处理的大小。nn.DataParallel 将批处理均匀地分配到多个 GPU。假设你有 4 个 GPU，批处理总大小为 32；然后，每个 GPU 将获得包含 8 个样本的块。但问题是，尽管所有的主 GPU 都可以轻松地将这些批处理放入对应的 VRAM 中，但主 GPU 必须分配额外的空间来容纳 32 个批处理大小，以用于其他卡的输出。

对于这种不均衡的 GPU 使用率，有两种解决方案：

1. 在训练期间继续在前向推导内使用 nn.DataParallel 计算损耗。在这种情况下。za 不会将密集的预测掩码返回给主 GPU，而只会返回单个标量损失；
2. 使用分布式训练，也称为 nn.DistributedDataParallel。借助分布式训练的另一个好处是可以看到 GPU 实现 100% 负载。

如果想了解更多，可以看看这三篇文章：

• https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
• https://medium.com/@theaccelerators/learn-pytorch-multi-gpu-properly-3eb976c030ee
• https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2

建议 5: 如果你拥有两个及以上的 GPU

能节省多少时间很大程度上取决于你的方案，我观察到，在 4x1080Ti 上训练图像分类 pipeline 时，大概可以节约 20% 的时间。另外值得一提的是，你也可以用 nn.DataParallel 和 nn.DistributedDataParallel 来进行推断。

关于自定义损失函数

编写自定义损失函数是一项很有趣的练习，我建议大家都不时尝试一下。提到这种逻辑复杂的损失函数，你要牢记一件事：它们都在 CUDA 上运行，你应该会写「CUDA-efficient」代码。「CUDA-efficient」意味着「没有 Python 控制流」。在 CPU 和 GPU 之间来回切换，访问 GPU 张量的个别值也可以完成这些工作，但是性能表现会很差。

前段时间，我实现了一个自定义余弦嵌入损失函数，是从《Segmenting and tracking cell instances with cosine embeddings and recurrent hourglass networks》这篇论文中来的，从文本形式上看它非常简单，但实现起来却有些复杂。

我编写的第一个简单实现的时候，（除了 bug 之外）花了几分钟来计算单个批的损失值。为了分析 CUDA 瓶颈，PyTorch 提供了一个非常方便的内置分析器，非常简单好用，提供了解决代码瓶颈的所有信息：

def test_loss_profiling():
    loss = nn.BCEWithLogitsLoss()
    with torch.autograd.profiler.profile(use_cuda=True) as prof:
        input = torch.randn((8, 1, 128, 128)).cuda()
        input.requires_grad = True

        target = torch.randint(1, (8, 1, 128, 128)).cuda().float()

        for i in range(10):
            l = loss(input, target)
            l.backward()
    print(prof.key_averages().table(sort_by="self_cpu_time_total"))

建议 9: 如果设计自定义模块和损失——配置并测试他们

在对最初的实现进行性能分析之后，就能够提速 100 倍。关于在 PyTorch 中编写高效张量表达式的更多信息，将在 Efficient PyTorch — Part 2 进行说明。

时间 VS 金钱

最后但非常重要的一点，有时候投资功能更强大的硬件，比优化代码可能更有价值。软件优化始终是结果无法确定的高风险之旅，升级 CPU、RAM、GPU 或者同时升级以上硬件可能会更有效果。金钱和时间都是资源，二者的均衡利用是成功的关键。

通过硬件升级可以更轻松地解决某些瓶颈。

写在最后

懂得充分利用日常工具是提高熟练度的关键，尽量不要制造「捷径」，如果遇到不清楚的地方，请深入挖掘，总有机会发现新知识。正所谓「每日一省」：问问自己，我的代码还能改进吗？这种精益求精的信念和其他技能一样，都是计算机工程师之路的必备品。

原文链接：https://towardsdatascience.com/efficient-pytorch-part-1-fe40ed5db76c