Pytorch - 分布式训练极简体验

本文通过一个简单的线性变换模型,演示如何在PyTorch中进行单机双卡的分布式训练。内容涵盖初始化通信模块、模型创建、损失函数与优化器、计算与梯度更新的详细步骤,旨在通过可执行的代码帮助理解分布式训练的实现过程。
摘要由CSDN通过智能技术生成

由于工作需要,最近在补充分布式训练方面的知识。经过一番理论学习后仍觉得意犹未尽,很多知识点无法准确get到(例如:分布式原语scatter、all reduce等代码层面应该是什么样的,ring all reduce 算法在梯度同步时是怎么使用的,parameter server参数是如何部分更新的)。

著名物理学家,诺贝尔奖得主Richard Feynman办公室的黑板上写了:"What I cannot create, I do not understand."。在程序员界也经常有"show me the code"的口号。 因此,我打算写一系列的分布式训练的文章,将以往抽象的分布式训练的概念以代码的形式展现出来,并保证每个代码可执行、可验证、可复现,并贡献出来源码让大家相互交流。

经过调研发现pytorch对于分布式训练做好很好的抽象且接口完善,因此本系列文章将以pytorch为主要框架进行,文章中的例子很多都来自pytorch的文档,并在此基础上进行了调试和扩充。

最后,由于分布式训练的理论介绍网络上已经很多了,理论部分的介绍不会是本系列文章的重点,我会将重点放在代码层面的介绍上面。

1 任务介绍

通过实现一个线性变换模型 y=xA^t+b 的单机2卡分布式训练任务,来初步体验pytorch中DistributeDataparallel的使用。本文主要参考pytorch tutorial中的介绍。

2 整体流程

代码编写流程如下:

  • 初始化pytorch分布式训练通信模块;
  • 创建模型(这里包括本地模型和分布式模型)
  • 创建损失函数和优化器
  • 计算(forward 和backward)和梯度更新
  • 多任务启动

3 初始化通信模块

pytorch中分布式通信模块为torch.distributed

本例中初始化代码为:

  • 通过环境变量MASTER_ADDRMASTER_PORT设置rank0的IP和PORT信息,rank0的作用相当于是协调节点,需要其他所有节点知道其访问地址;
  • <
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值