torch.DistributedDataParallel复盘

最新推荐文章于 2024-03-15 20:59:38 发布

zkxhlbt

最新推荐文章于 2024-03-15 20:59:38 发布

阅读量262

点赞数 1

分类专栏： CV 文章标签： pytorch

本文链接：https://blog.csdn.net/zkxhlbt/article/details/119318972

版权

14 篇文章 0 订阅

订阅专栏

用多GPU跑torch程序，DDP是最好的工具。在研究过程中基本是单机多GPU而不是多机多GPU，此处对单机多GPU的情况做一个总结。
import torch.DistributedDataParallel as DDP
import torch.distributed as dist
用法：

由于是单机多GPU，所以rank == local_rank. 使用argparse从外部获取–local_rank后，加入dist.init_process_group(backend=‘nccl’) 这一句话来初始化。
数据集方面。使用sampler = th.utils.data.distributed.DistributedSampler(training_set) 来定义training set 的sampler,这样可以保证在训练时，不同GPU使用不同（没有交集）的batch.注意！！！在测试的时候要想清楚测试数据的顺序是不是有影响。例如今天写的代码，要使用前四个测试数据来画图，这个时候shuffle要设为false, sampler也不用定义，一定不能惯性思维。batch_size的设置也要格外留意。
网络方面，要在.cuda()后加上model = DDP(model, device_ids=[dist.get_rank()], output_device=dist.get_rank()) # 再次强调，在单机多卡的情况下，rank == local_rank，所以在这里直接使用rank就行
.cuda的问题。在这里，模型和数据不能简单地input.cuda(), model.cuda()，而要在cuda()中指定device,否则会有数据和模型在不同 GPU上的情况。正确的做法是 input.cuda(torch.device(“cuda”, dist.get_rank())).
学习率要随着batch_size的变化而变化。batch_size增大几倍，学习率就要相应增大几倍。