计图MPI分布式多卡
计图分布式基于MPI（Message Passing Interface），主要阐述使用计图MPI，进行多卡和分布式训练。目前计图分布式处于测试阶段。
计图MPI安装
计图依赖OpenMPI，用户可以使用如下命令安装OpenMPI：
sudo apt install openmpi-bin openmpi-common libopenmpi-dev
计图会自动检测环境变量中是否包含mpicc，如果计图成功的检测到了mpicc，输出如下信息：
[i 0502 14:09:55.758481 24 init.py:203] Found mpicc(1.10.2) at /usr/bin/mpicc
如果计图没有在环境变量中找到mpi，用户也可以手动指定mpicc的路径告诉计图，添加环境变量即可：export mpicc_path=/you/mpicc/path
OpenMPI安装完成以后，用户无需修改代码，需要做的仅仅是修改启动命令行，计图就会用数据并行的方式，自动完成并行操作。

单卡训练代码

python3.7 -m jittor.test.test_resnet

分布式多卡训练代码

mpirun -np 4 python3.7 -m jittor.test.test_resnet

指定特定显卡的多卡训练代码

CUDA_VISIBLE_DEVICES=“2,3” mpirun -np 2 python3.7 -m jittor.test.test_resnet
便捷性的背后，计图的分布式算子的支撑，计图支持的mpi算子后端会使用nccl进行进一步的加速。计图所有分布式算法的开发，均在Python前端完成，让分布式算法的灵活度增强，开发分布式算法的难度也大大降低。
基于这些mpi算子接口，研发团队已经集成了如下三种分布式相关的算法：
• 分布式数据并行加载
• 分布式优化器
• 分布式同步批归一化层
用户在使用MPI进行分布式训练时，计图内部的Dataset类会自动并行分发数据，需要注意的是Dataset类中设置的Batch size是所有节点的batch size之和，也就是总batch size，不是单个节点接收到的batch size。
MPI接口
目前MPI开放接口如下：
• jt.mpi: 计图的MPI模块，当计图不在MPI环境下时，jt.mpi == None，用户可以用这个判断是否在mpi环境下。
• jt.Module.mpi_param_broadcast(root=0): 将模块的参数从root节点广播给其他节点。
• jt.mpi.mpi_reduce(x, op=‘add’, root=0): 将所有节点的变量x使用算子op，reduce到root节点。如果op是’add’或者’sum’，该接口会把所有变量求和，如果op是’mean’，该接口会取均值。

• jt.mpi.mpi_broadcast(x, root=0): 将变量x从root节点广播到所有节点。

• jt.mpi.mpi_all_reduce(x, op=‘add’): 将所有节点的变量x使用一起reduce，并且吧reduce的结果再次广播到所有节点。如果op是’add’或者’sum’，该接口会把所有变量求和，如果op是’mean’，该接口会取均值。

实例：MPI实现分布式同步批归一化层
下面的代码是使用计图实现分布式同步批，归一化层的实例代码，在原来批归一化层的基础上，只需增加三行代码，就可以实现分布式的batch norm，添加的代码如下：

将均值和方差，通过all reduce同步到所有节点

if self.sync and jt.mpi:
xmean = xmean.mpi_all_reduce(“mean”)
x2mean = x2mean.mpi_all_reduce(“mean”)
注：计图内部已经实现了同步的批归一化层，用户不需要自己实现
分布式同步批归一化层的完整代码：
class BatchNorm(Module):
def init(self, num_features, eps=1e-5, momentum=0.1, affine=None, is_train=True, sync=True):
assert affine == None

    self.sync = syncself.num_features = num_featuresself.is_train = is_trainself.eps = epsself.momentum = momentumself.weight = init.constant((num_features,), "float32", 1.0)self.bias = init.constant((num_features,), "float32", 0.0)self.running_mean = init.constant((num_features,), "float32", 0.0).stop_grad()self.running_var = init.constant((num_features,), "float32", 1.0).stop_grad()def execute(self, x):if self.is_train:xmean = jt.mean(x, dims=[0,2,3], keepdims=1)x2mean = jt.mean(x*x, dims=[0,2,3], keepdims=1)# 将均值和方差，通过all reduce同步到所有节点if self.sync and jt.mpi:xmean = xmean.mpi_all_reduce("mean")x2mean = x2mean.mpi_all_reduce("mean")xvar = x2mean-xmean*xmeannorm_x = (x-xmean)/jt.sqrt(xvar+self.eps)self.running_mean += (xmean.sum([0,2,3])-self.running_mean)*self.momentumself.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentumelse:running_mean = self.running_mean.broadcast(x, [0,2,3])running_var = self.running_var.broadcast(x, [0,2,3])norm_x = (x-running_mean)/jt.sqrt(running_var+self.eps)w = self.weight.broadcast(x, [0,2,3])b = self.bias.broadcast(x, [0,2,3])return norm_x * w + b

计图MPI分布式多卡相关推荐

计图(Jittor) 1.1版本：新增骨干网络、JIT功能升级、支持多卡训练
计图(Jittor) 1.1版本:新增骨干网络.JIT功能升级.支持多卡训练深度学习框架-计图(Jittor),Jittor的新版本V1.1上线了.主要变化包括: • 增加了大量骨干网络的支持,增强 ...
清华「计图」、旷视「天元」纷纷重磅开源，国产深度学习框架迎来高光时刻...
来源:CSDN 本文约3141字,建议阅读7分钟. 本文介绍清华开源的深度学习框架 Jittor(计图) 引起了广泛关注,支撑旷视诸多业务和战绩的深度学习框架 MegEngine(天元) 也将在近日开 ...
ResNet也能用在3D模型上了，清华「计图」团队新研究已开源
鱼羊发自凹非寺量子位报道 | 公众号 QbitAI 用AI处理二维图像,离不开卷积神经网络(CNN)这个地基. 不过,面对三维模型,CNN就没有那么得劲了. 主要原因是,3D模型通常采用网格数 ...
清华自研深度学习框架「计图」开源！多项任务性能超过PyTorch
2020-03-20 11:42:07 乾明发自凹非寺量子位报道 | 公众号 QbitAI AI框架,又来重磅中国玩家. 刚刚,清华自研的深度学习框架,正式对外开源."贵系&quo ...
【Redis | 黑马点评 + 思维导图】分布式锁
文章目录分布式锁的基本原理和实现方式对比 Redis分布式锁的实现核心思路分布式锁的初级实现 Redis分布式锁误删情况说明解决Redis分布式锁误删问题分布式锁的原子性问题 Lua脚本解决多 ...
清华大学胡事民是计算机视觉,ResNet也能用在3D模型上了，清华「计图」团队新研究已开源...
鱼羊发自凹非寺量子位报道 | 公众号 QbitAI 用AI处理二维图像,离不开卷积神经网络(CNN)这个地基. 不过,面对三维模型,CNN就没有那么得劲了. 主要原因是,3D模型通常采用网格数 ...
清华大学胡事民是计算机视觉,ResNet也能用在3D模型上了，清华“计图”团队新研究已开源...
鱼羊发自凹非寺量子位报道 | 公众号 QbitAI 用AI处理二维图像,离不开卷积神经网络(CNN)这个地基. 不过,面对三维模型,CNN就没有那么得劲了. 主要原因是,3D模型通常采用网格数 ...
盘点国内那些深度学习框架：清华计图Jittor、腾讯优图NCNN、百度飞桨PaddlePaddle、阿里X-DeepLearning
1.清华计图Jittor 清华大学开发了一个名为计图(Jittor)的深度学习框架. 计图(Jittor:Just in Time)是一个采用元算子表达神经网络计算单元.完全基于动态编译(Just-i ...
对标Pytorch，清华团队推出自研AI框架“计图”
「AI技术生态论」人物访谈栏目是CSDN发起的百万人学AI倡议下的重要组成部分.通过对AI生态专家.创业者.行业KOL的访谈,反映其对于行业的思考.未来趋势的判断.技术的实践,以及成长的经历. 20 ...

计图MPI分布式多卡

单卡训练代码

分布式多卡训练代码

指定特定显卡的多卡训练代码

将均值和方差，通过all reduce同步到所有节点

计图MPI分布式多卡相关推荐

最新文章

热门文章