使用DistrbutedDataParallel时，nvdiai-smi显示每个进程都占用GPU:0

使用mp.spawn(main, nprocs=args.num_gpus, args=(cfg, args,))和torch.nn.distributed.DistrbutedDataParallel时出现显卡在

用torch.load加载预训练模型或恢复训练时必须指定放在哪个GPU上，否则默认是每个进程都会在GPU:0 上占一块空间，用来存放恢复数据，而且很大。nvidia-smi的结果中可以看到0号显卡上会有每个进程的编号。

 

nvidia-smi结果

错误做法：

checkpoint = torch.load(pretrain_path)

正确做法：

checkpoint = torch.load(pretrain_path, map_location= lambda storage, loc: storage.cuda(args.local_rank))

args.local_rank是计算出来当前显卡在节点内的rank。