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（路线图上的顶点）

推荐阅读

国产小众浏览器因屏蔽视频广告，被索赔100万（后续）

年轻人“不讲武德”：因看黄片上瘾，把网站和786名女主播起诉了

中国联通官网被发现含木马脚本，可向用户推广色情APP

张一鸣：每个逆袭的年轻人，都具备的底层能力

关于程序员大白

程序员大白是一群哈工大，东北大学，西湖大学和上海交通大学的硕士博士运营维护的号，大家乐于分享高质量文章，喜欢总结知识，欢迎关注[程序员大白]，大家一起学习进步！