PyTorch多GPU并行训练方法及问题整理

os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

model = nn.DataParallel(model)model = model.cuda()

inputs = inputs.cuda()labels = labels.cuda()

model = Model(input_size, output_size)if torch.cuda.device_count() > 1:    print("Let's use", torch.cuda.device_count(), "GPUs!")    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs    model = nn.DataParallel(model)
model.to(device)

class DataParallel(Module):    def __init__(self, module, device_ids=None, output_device=None, dim=0):        super(DataParallel, self).__init__()
        if not torch.cuda.is_available():            self.module = module            self.device_ids = []            return
        if device_ids is None:            device_ids = list(range(torch.cuda.device_count()))        if output_device is None:            output_device = device_ids[0]

class BalancedDataParallel(DataParallel):    def __init__(self, gpu0_bsz, *args, **kwargs):        self.gpu0_bsz = gpu0_bsz        super().__init__(*args, **kwargs)    ...

my_net = MyNet()my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()

batch_szie = 8gpu0_bsz = 2acc_grad = 1my_net = MyNet()my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()

batch_szie = 16gpu0_bsz = 4acc_grad = 2my_net = MyNet()my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()

torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)

torch.distributed.init_process_group(backend="nccl")model = DistributedDataParallel(model) # device_ids will include all GPU devices by default

python -m torch.distributed.launch main.py

model = model.cuda()model = nn.parallel.DistributedDataParallel(model)

backend: 后端, 实际上是多个机器之间交换数据的协议init_method: 机器之间交换数据, 需要指定一个主节点, 而这个参数就是指定主节点的world_size: 介绍都是说是进程, 实际就是机器的个数, 例如两台机器一起训练的话, world_size就设置为2rank: 区分主节点和从节点的, 主节点为0, 剩余的为了1-(N-1), N为要使用的机器的数量, 也就是world_size

import os# 以下二选一, 第一个是使用gloo后端需要设置的, 第二个是使用nccl需要设置的os.environ['GLOO_SOCKET_IFNAME'] = 'eth0'os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'

import torch.distributed as dist
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',                        rank=rank, world_size=world_size)

import torch.distributed as dist
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',                        rank=rank, world_size=world_size)

python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 0 -ws 2

python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 1 -ws 2

print("Initialize Dataloaders...")# Define the transform for the data. Notice, we must resize to 224x224 with this dataset and model.transform = transforms.Compose(    [transforms.Resize(224),     transforms.ToTensor(),     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# Initialize Datasets. STL10 will automatically download if not presenttrainset = datasets.STL10(root='./data', split='train', download=True, transform=transform)valset = datasets.STL10(root='./data', split='test', download=True, transform=transform)
# Create DistributedSampler to handle distributing the dataset across nodes when training# This can only be called after torch.distributed.init_process_group is called# 这一句就是和平时使用有点不一样的地方train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
# Create the Dataloaders to feed data to the training and validation stepstrain_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=(train_sampler is None), num_workers=workers, pin_memory=False, sampler=train_sampler)val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=False)

model = model.cuda()model = nn.parallel.DistributedDataParallel(model)

def demo_checkpoint(rank, world_size):    setup(rank, world_size)
    # setup devices for this process, rank 1 uses GPUs [0, 1, 2, 3] and    # rank 2 uses GPUs [4, 5, 6, 7].    n = torch.cuda.device_count() // world_size    device_ids = list(range(rank * n, (rank + 1) * n))
    model = ToyModel().to(device_ids[0])    # output_device defaults to device_ids[0]    ddp_model = DDP(model, device_ids=device_ids)
    loss_fn = nn.MSELoss()    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    CHECKPOINT_PATH = tempfile.gettempdir() + "/model.checkpoint"    if rank == 0:        # All processes should see same parameters as they all start from same        # random parameters and gradients are synchronized in backward passes.        # Therefore, saving it in one process is sufficient.        torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)
    # Use a barrier() to make sure that process 1 loads the model after process    # 0 saves it.    dist.barrier()    # configure map_location properly    rank0_devices = [x - rank * len(device_ids) for x in device_ids]    device_pairs = zip(rank0_devices, device_ids)    map_location = {'cuda:%d' % x: 'cuda:%d' % y for x, y in device_pairs}    ddp_model.load_state_dict(        torch.load(CHECKPOINT_PATH, map_location=map_location))
    optimizer.zero_grad()    outputs = ddp_model(torch.randn(20, 10))    labels = torch.randn(20, 5).to(device_ids[0])    loss_fn = nn.MSELoss()    loss_fn(outputs, labels).backward()    optimizer.step()
    # Use a barrier() to make sure that all processes have finished reading the    # checkpoint    dist.barrier()
    if rank == 0:        os.remove(CHECKPOINT_PATH)
    cleanup()