pytorch里DataParallel 和 DistributedParallel

参考博客：
http://aiuai.cn/aifarm1340.html
https://yangkky.github.io/2019/07/08/distributed-pytorch-tutorial.html

pytorch里DataParallel 和 DistributedParallel 的使用说明

1 DataParallel模式
以一张卡作为主卡，向其他卡分发数据，汇总loss和梯度的方式来训练，速度有提升
# 参考这个解释： https://blog.csdn.net/weixin_40087578/article/details/87186613
实现方式：

# 数据加载 不用改
# 模型部分修改
model=model.cuda()
model=torch.nn.DaraParallel(model)
# 前向和反向也不用改

 

2 DistributedParallel模式
多张卡独立进程进行加载数据和前方反向传播，分布式传递数据，速度很快，推荐这种
参考这个解释：https://www.cnblogs.com/yh-blog/p/12877922.html
实现方式(只适用于单机多卡)：

# 初始化 增加
import torch.distributed as dist
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=0, help='node rank for distributed parallel')
parser.add_argument('--distributed', type=int, default=0, help='distributed mode')
args = parser.parse_args()
assert torch.distributed.is_nccl_available()
torch.cuda.set_device(args.local_rank)
device_num = torch.cuda.device_count()
distributed_mode = device_num >= 2 and args.distributed
if distributed_mode:
dist.init_process_group('nccl', world_size=device_num, rank=args.local_rank)
rank = dist.get_rank()
num_rep = dist.get_world_size()
print(rank, num_rep)
print('torch distributed work is inited.')
# 数据加载 修改，多卡会按照自己的进程id进行分片加载
ds = MyDataset(root, transform)
_sampler = dd.distributed.DistributedSampler(ds,num_replicas=gpu_num)
dl = DataLoader(ds,sampler=_sampler,shuffle=_sampler is None,batch_size=128)

# 模型部分修改
model=model.cuda()
# 注意，BN层一定要转化成SBN，
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = nn.parallel.DistributedDataParallel(model)
# 使用tensorboard要注意，由于是多个进程一起跑，会出现同时读写writer，会报错，所以只在一个rank上写。其他rank不写日志
if dist.get_rank()==0：
　　写log； 保存summary到writer
else：
　　跳过
# 前向和反向也不用改
启动方式：
　　CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nnodes=1 --nproc_per_node=2 --master_port=29501 main.py --distributed=1