torch：DistributedDataParallel分布式单机多GPUS训练

• 前言
• 一、快速在代码里加入DDP

pytorch官方链接
torch官网对于多卡训练有较详细讲解，并有多种实现方式。这里也不再班门弄斧，只记录DDP分布式训练的方法应用，来帮助没经验的人快速用起来。

使用过程中总结大概需要以下几步：
1、设置local_rank及device(启动器开启后，会将当前进程的（其实就是 GPU的）index 通过参数传递给python环境local_rank参数中，os.environ[“LOCAL_RANK”]也是可以获取到)
2、初始化分布式环境、进程间通信
3、DDP包装model
4、采样器包装dataset
5、数据输入to.(device)或者.cuda()
6、用torch.distributed.launch启动（其他方式不作分享）

#################################################
    # pytorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：
    parser = argparse.ArgumentParser()
	parser.add_argument("--local_rank", type=int, default=-1)
	args = parser.parse_args()
	# 方式一：args.local_rank 方式二：int(os.environ['LOCAL_RANK']) 
	# 以上两种方式都可以获取，当前GPU或者说进程的ID，说白了就是当前进程是在哪张GPU显卡
	# args.local_rank == int(os.environ['LOCAL_RANK'])
    # set the device
    args.device = torch.cuda.device(args.local_rank)
    dist.init_process_group(backend='nccl', init_method='env://')
    torch.cuda.set_device(args.local_rank)

    # set the seed for all GPUs (also make sure to set the seed for random, numpy, etc.)
    torch.cuda.manual_seed_all(SEED)
###############################################
    # initialize your model 
    model = YourModel()
    # send your model to GPU
    model = model.to(device)
    # initialize distributed data parallel (DDP)
    model = DDP(
        model,
        device_ids=[args.local_rank],
        output_device=args.local_rank
    )
###########################################
    # initialize your dataset
    dataset = YourDataset()
    # initialize the DistributedSampler
    sampler = DistributedSampler(dataset)

    # initialize the dataloader
    dataloader = DataLoader(
        dataset=dataset,
        sampler=sampler,
        batch_size=BATCH_SIZE,
        broadcast_buffers=False,
        find_unused_parameters=True
    )
################################################
    # start your training!
    for epoch in range(NUM_EPOCHS):
        # put model in train mode
        model.train()
        # forward
        # data = data.cuda(non_blocking=True)
        inputs = inputs.to(device)
	    labels = labels.to(device)
	    outputs = model(inputs)
	    loss = criterion(outputs, labels)
	    # backward
	    optimizer.zero_grad()
	    loss.backward()
	    optimizer.step()

        # let all processes sync up before starting with a new epoch of training
        dist.barrier()

#!/bin/bash

# this example uses a single node (`NUM_NODES=1`) w/ 4 GPUs (`NUM_GPUS_PER_NODE=4`)
export NUM_NODES=1
export NUM_GPUS_PER_NODE=4
export NODE_RANK=0
export WORLD_SIZE=$(($NUM_NODES * $NUM_GPUS_PER_NODE))

# launch your script w/ `torch.distributed.launch`
python -m torch.distributed.launch 
    --nproc_per_node=$NUM_GPUS_PER_NODE 
    --nnodes=$NUM_NODES 
    --node_rank $NODE_RANK 
    ddp_example.py