Pytorch中多GPU训练指北

如果我们机子中有很多显卡(例如我们有八张显卡)，但我们只想使用0、1、2号显卡，那么我们可以：

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

另一种方法DistributedParallel

python -m torch.distributed.launch --nproc_per_node=你的GPU数量
               YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
               arguments of your training script)

为此，我们的YOUR_TRAINING_SCRIPT.py也就是我们的训练代码中这样写(省略多余代码，只保留核心代码)：

import torch.distributed as dist
# 这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
parser.add_argument("--local_rank", type=int, default=0) 

def synchronize():
    """
    Helper function to synchronize (barrier) among all processes when
    using distributed training
    """
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return
    world_size = dist.get_world_size()
    if world_size == 1:
        return
    dist.barrier()


## WORLD_SIZE 由torch.distributed.launch.py产生 具体数值为 nproc_per_node*node(主机数，这里为1)
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

is_distributed = num_gpus > 1

if is_distributed:
    torch.cuda.set_device(args.local_rank)  # 这里设定每一个进程使用的GPU是一定的
    torch.distributed.init_process_group(
        backend="nccl", init_method="env://"
    )
    synchronize()

# 将模型移至到DistributedDataParallel中，此时就可以进行训练了
if is_distributed:
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[args.local_rank], output_device=args.local_rank,
        # this should be removed if we update BatchNorm stats
        broadcast_buffers=True,
    )

# 注意，在测试的时候需要执行 model = model.module