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



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

极市平台

2021-04-22

↑ 点击蓝字关注极市平台

作者丨力元

来源丨机器之心

编辑丨极市平台

极市导读

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

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

如果觉得有用，就请分享到朋友圈吧！△点击卡片关注极市平台，获取最新CV干货

推荐阅读

硬盘告急！没想到矿工开始用硬盘挖币，电商库存几乎被一扫而空

2021-04-21

七天近千星！哈佛小哥这个github仓库从零开始教你计算机图形学

2021-04-20

一位中科院博士的毕业论文致谢：求学22载，计算机终成一生的事业与希望

2021-04-19

# CV技术社群邀请函 #

△长按添加极市小助手

添加极市小助手微信（ID : cvmart2）

备注：姓名-学校/公司-研究方向-城市（如：小极-北大-目标检测-深圳）

即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群

每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~

觉得有用麻烦给个在看啦~

【声明】内容源于网络

极市平台

为计算机视觉开发者提供全流程算法开发训练平台，以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。

内容 8155

粉丝 0

极市平台为计算机视觉开发者提供全流程算法开发训练平台，以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。

总阅读7.6k

粉丝0

内容8.2k