之前看了很多例子,始终不明白rank是怎么导入进去的,现在跑通了,特此记录。具体代码和参考差不多。这里是单机多卡的使用经验,多机多卡的可以不用看了。下面的代码我给出的参考里都有,这里只记录跑的时候一些报错并更正的地方。
举例:我现在有两张卡,希望在两张卡上进行分布式训练。
rank:目前所用的卡
world_size :你有多少个进程,通常是一张卡一个进程
world_size = torch.cuda.device_count()
1.dist.init_process_group初始化
这步踩过的坑就是init_method不知道怎么设置。
rank怎么来的后面说。这一步就是设定一个组,这个组里面设定你有几个进程(world size),现在是卡几(rank)。让pycharm知道你要跑几个进程,包装在组内,进行通讯这样模型参数会自己同步,不需要额外操作了。
import os
import torch.distributed as dist
def ddp_setup(rank,world_size):
os.environ['MASTER_ADDR'] = 'localhost' #rank0 对应的地址
os.environ['MASTER_PORT'] = '12355' #任何空闲的端口
dist.init_process_group(
backend='nccl', #nccl Gloo #nvidia显卡的选择nccl
world_size=world_size,
init_method='env://',
rank=rank
) #初始化默认的分布进程组
dist.barrier() #等到每块GPU运行到这再继续往下走
2.spawn启动(rank怎么来的)
rank是自动分配的。怎么分配呢?这里用的是spawn也就一行代码。
import torch.multiprocessing as mp
def main (rank:int,world_size:int,args):
pass
#你的训练代码 主函数
mp.spawn(main,args=(args.world_size,args), nprocs=args.world_size)
注意,我调用spawn的时候,没有输入main的其中一个参数rank。这行代表rank由代码自动分配。将你的代码复制两份在两张卡上同时跑,你可以print(rank),会发现输出 0 1。两份代码并行跑。
另外,nprocs=args.world_size。如果不这么写,我的代码会卡死在这,即不报错,也不停止。
参考:
1.多卡训练代码PyTorch单机多卡分布式训练(源代码讲解)_pytorch单机多卡训练_Wukongggg的博客-CSDN博客2
2.多卡训练B站教程:pytorch多GPU并行训练教程_哔哩哔哩_bilibili