Hugging Face发布PyTorch新库「Accelerate」：适用于多GPU、TPU、混合精度训练
机器之心报道
多数 PyTorch 高级库都支持分布式训练和混合精度训练，但是它们引入的抽象化往往需要用户学习新的 API 来定制训练循环。许多 PyTorch 用户希望完全控制自己的训练循环，但不想编写和维护训练所需的样板代码。Hugging Face 最近发布的新库 Accelerate 解决了这个问题。
import torchimport torch.nn.functional as Ffrom datasets import load_dataset+ from accelerate import Accelerator+ accelerator = Accelerator()- device = 'cpu'+ device = accelerator.devicemodel = 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()
import torchimport torch.nn.functional as Ffrom 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()
accelerate configaccelerate launch my_script.py --args_to_my_scriptaccelerator = Accelerator()model, optim, data = accelerator.prepare(model, optim, data)accelerator.backward(loss)CPU
单 GPU
单一节点多 GPU
多节点多 GPU
TPU
带有本地 AMP 的 FP16（路线图上的顶点）
建新·见智 —— 2021亚马逊云科技 AI 在线大会
4月22日 14:00 - 18:00
《建新 · 见智——2021 亚马逊云科技 AI 在线大会》由亚马逊云科技全球人工智能技术副总裁及杰出科学家 Alex Smola、亚马逊云科技大中华区产品部总经理顾凡领衔，40多位重磅嘉宾将在主题演讲及6大分会场上为你深度剖析亚马逊云科技创新文化，揭秘 AI/ML 如何帮助企业加速创新。
分会场一：亚马逊机器学习实践揭秘
分会场二：人工智能赋能企业数字化转型
分会场三：大规模机器学习实现之道
分会场四：AI 服务助力互联网快速创新
分会场五：开源开放与前沿趋
分会场六：合作共赢的智能生态
6大分会场，你对哪个主题更感兴趣？
识别二维码或点击阅读原文，免费报名看直播。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：content@jiqizhixin.com
登录查看更多
相关内容
专知会员服务
26+阅读 · 2020年3月26日
专知会员服务
61+阅读 · 2019年12月29日
Arxiv
0+阅读 · 2021年6月23日
Arxiv
0+阅读 · 2021年6月20日
相关VIP内容
专知会员服务
26+阅读 · 2020年3月26日
专知会员服务
61+阅读 · 2019年12月29日
相关资讯
相关论文
Arxiv
0+阅读 · 2021年6月23日
Arxiv
0+阅读 · 2021年6月20日
微信扫码咨询专知VIP会员