pytorch DistributedDataParallel多卡并行训练

pytorch DistributedDataParallel多卡并行训练

Pytorch 中最简单的并行计算方式是 nn.DataParallel。
DataParallel 使用单进程控制将模型和数据加载到多个 GPU 中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。

但是DataParallel的缺点十分明显，各卡之间的负载不均衡，主卡的负载过大。运行时间大概是distributed 的四倍。
所以，下面我们介绍使用distributed 的计算方式。

先设定下local_rank，这是很重要的参数

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank