pytorch分布式训练理解

单机多卡训练

配置每个进程的gpu

• ddp模式

local_rank = torch.distributed.get_rank()	# 获取当前进程在所有进程中的编号
torch.cuda.set_device(local_rank)	# 与下面的等效使用CUDA_VISIBLE_DEVICES=local_rank，即指定当前进程使用的gpu编号
device = torch.device("cuda", local_rank) ## 表示 torch.Tensor 分配到指定编号的设备
device = torch.device("cuda:1") #如果local_rank=1,与上述等效
device = torch.device("cuda")  # 默认为当前设备

• horovod模式

hvd.init()
torch.cuda.set_device(hvd.local_rank())
#此处没有申明device，这样，在将torch.Tensor分配到gpu时，可以采用torch.cuda()

cuda其他用法

#cuda是否可用
print(torch.cuda.is_available())
##返回gpu数量
print(torch.cuda.device_count())
#返回当前设备索引；
print(torch.cuda.current_device())