PyTorch 口袋参考手册

第 6 章 PyTorch 加速和优化 PyTorch 加速和优化

本作品已使用人工智能进行翻译。欢迎您提供反馈和意见：[e-mail]

在前面的章节中，你学习了如何使用 PyTorch 的内置功能，并通过创建自己的深度学习自定义组件来扩展这些功能。这样做可以让你快速设计新模型和训练它们的算法。

然而，在处理超大数据集或更复杂的模型时，在单个 CPU 或 GPU 上训练模型可能会耗费很长时间--可能需要几天甚至几周才能得到初步结果。较长的训练时间可能会令人沮丧，尤其是当您想使用不同的超参数配置进行多次实验时。

在本章中，我们将探索使用 PyTorch 加速和优化模型开发的最新技术。首先，我们将了解使用张量处理单元（TPU）而不是 GPU 设备的情况，并考虑使用 TPU 可以提高性能的实例。接下来，我将向您展示如何使用 PyTorch 的内置功能进行并行处理和分布式训练。这将为在多个 GPU 和多台机器上训练模型提供快速参考，以便在有更多硬件资源可用时快速扩展训练。在探索加速训练的方法之后，我们将了解如何使用超参数调整、量化和剪枝等高级技术来优化模型。

本章还将提供参考代码，使入门更加容易，并提供我们使用过的关键软件包和库的参考。创建模型和训练循环后，您可以返回本章了解如何加速和优化训练过程的技巧。

让我们从探索如何在 TPU 上运行模型开始。

TPU 上的 PyTorch

随着 Deep Learning 和人工智能的部署越来越多，各家公司都在开发定制硬件芯片或 ASIC，旨在优化硬件中的模型性能。 谷歌开发了自己的 ASIC，用于 NN 加速，名为 TPU。由于 TPU 是专为 NN 设计的，因此它没有 GPU 的一些缺点，而 GPU 是专为图形处理设计的。谷歌的 TPU 现已作为 Google Cloud TPU 的一部分供您使用。 您还可以使用 TPU 运行 Google Colab。

在前面的章节中，我向您展示了如何使用 GPU 测试和训练深度模型。如果以下条件适用于您的使用案例，您应该继续使用 CPU 和 GPU 进行训练：

• 你们有小批量的小型或中型机型。

• 您的模型不需要很长时间就能训练出来。

• 数据进出是主要瓶颈。

• 您的计算经常使用分支，或者主要以元素为单位进行计算，或者使用稀疏内存访问。

• 您需要使用高精度。热塑性聚氨酯不适合双面胶。

另一方面，使用 TPU 而不是 GPU 进行训练有几个原因。TPU 在执行密集向量和矩阵计算时速度非常快。它们针对特定的工作负载进行了优化。当以下情况适用于您的使用案例时，您应强烈考虑使用 TPU：

• 您的模型以矩阵计算为主。

• 您的模型训练时间较长。

• 您希望在 TPU 上对整个训练循环进行多次迭代。

在 TPU 上运行与在 CPU 或 GPU 上运行非常相似。让我们在下面的代码中重温一下如何在 GPU 上训练模型：

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

model.to(device) 
for epoch in range(n_epochs):
  for data in trainloader:
    input, labels ...