Hugging Face发布PyTorch新库「Accelerate」：适用于多GPU、TPU、混合精度训练

来自：机器之心    作者：力元

多数 PyTorch 高级库都支持分布式训练和混合精度训练，但是它们引入的抽象化往往需要用户学习新的 API 来定制训练循环。许多 PyTorch 用户希望完全控制自己的训练循环，但不想编写和维护训练所需的样板代码。Hugging Face 最近发布的新库 Accelerate 解决了这个问题。

「Accelerate」提供了一个简单的 API，将与多 GPU 、 TPU 、 fp16 相关的样板代码抽离了出来，保持其余代码不变。PyTorch 用户无须使用不便控制和调整的抽象类或编写、维护样板代码，就可以直接上手多 GPU 或 TPU。

项目地址：https://github.com/huggingface/accelerate

通过将如下 5 行代码添加到原始的 PyTorch 训练循环中，脚本即可在本地以及任何分布式设置上运行。

 

import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator


+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.device


  model = torch.nn.Transformer().to(device)
  optim = torch.optim.Adam(model.parameters())


  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)


+ model, optim, data = accelerator.prepare(model, optim, data)


  model.train()
  for epoch in range(10):
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)


          optimizer.zero_grad()


          output = model(source)
          loss = F.cross_entropy(output, targets)


+         accelerator.backward(loss)
-         loss.backward()


          optimizer.step()

Accelerate 甚至可以通过处理设备的放置（需要对代码进行一些更改，但通常更安全）进一步简化训练循环，代码如下所示： 

import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator


+ accelerator = Accelerator()
- device = 'cpu'


+ model = torch.nn.Transformer()
- model = torch.nn.Transformer().to(device)
  optim = torch.optim.Adam(model.parameters())


  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)


+ model, optim, data = accelerator.prepare(model, optim, data)


  model.train()
  for epoch in range(10):
      for source, targets in data:
-         source = source.to(device)
-         targets = targets.to(device)


          optimizer.zero_grad()


          output = model(source)
          loss = F.cross_entropy(output, targets)


+         accelerator.backward(loss)
-         loss.backward()


          optimizer.step()

除了 API，Accelerate 还提供了一个 CLI 工具，方便启动脚本之前快速配置和测试训练环境，然后启动脚本。指令如下所示： 

accelerate config

accelerate launch my_script.py --args_to_my_script

如果不想自己编写训练循环， PyTorch 之上有许多可以替代 Accelerate 的高级库。

Accelerate 的运作原理

accelerator = Accelerator()

除了提供要使用的主要对象之外，此行还将从环境中分析分布式训练运行的类型并执行必要的初始化。用户可以通过将 cpu = True 或 fp16 = True 传递给此 init 来强制进行 CPU 训练或混合精度训练。这两个选项都可以使用脚本的启动器进行设置。

model, optim, data = accelerator.prepare(model, optim, data)

这是 API 的主体，将准备三种主要类型的对象：models (torch.nn.Module)、optimizers (torch.optim.Optimizer)、dataloaders (torch.data.dataloader.DataLoader)。

模型

模型的准备包括将其包装在适当的容器（例如 DistributedDataParallel）中，然后将其放置在适当的设备上。与普通分布式训练一样，进行保存或访问其特定的方法时，需要先通过 accelerator.unwrap_model（model）解开模型。

 

优化器

优化器也包装于一个为使混合精度训练运转执行必要操作的特殊容器里。如果状态 dict 是非空的或从检查点加载的，它会准确地处理状态 dict 的设备放置。

数据加载器

此库并不依赖于 DistributedSampler，它实际上可以与传递到数据加载器的采样器一起使用。数据加载器包装于仅在采样器中获取与当前进程相关的索引并将批次放入设备的容器中。

 

为此，Accelerate 提供了一种实用程序功能来同步更多 RNGs。该功能将在分布式训练期间运行的每个进程上同步随机数生成器。默认情况下，它仅同步采样器的生成器，因此在每个进程中数据扩充都将不同，但是随机改组将是相同的。

accelerator.backward(loss)

此行代码为向后传递添加了必要的步骤来提高混合精度，但对于其他集成则需要进行一些自定义。

Accelerate 支持的集成包括：

• CPU

• 单 GPU

• 单一节点多 GPU

• 多节点多 GPU

• TPU

• 带有本地 AMP 的 FP16（路线图上的顶点）

 End 

声明：部分内容来源于网络，仅供读者学术交流之目的，文章版权归原作者所有。如有不妥，请联系删除。

猜您喜欢：

等你着陆！【GAN生成对抗网络】知识星球！

超100篇！CVPR 2020最全GAN论文梳理汇总！

附下载 | 《Python进阶》中文版

附下载 | 经典《Think Python》中文版

附下载 | 《Pytorch模型训练实用教程》

附下载 | 最新2020李沐《动手学深度学习》

附下载 | 《可解释的机器学习》中文版

附下载 |《TensorFlow 2.0 深度学习算法实战》

附下载 | 超100篇！CVPR 2020最全GAN论文梳理汇总！

附下载 |《计算机视觉中的数学方法》分享