PyTorch多GPU并行訓(xùn)練方法及問(wèn)題整理

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: 后端, 實(shí)際上是多個(gè)機(jī)器之間交換數(shù)據(jù)的協(xié)議init_method: 機(jī)器之間交換數(shù)據(jù), 需要指定一個(gè)主節(jié)點(diǎn), 而這個(gè)參數(shù)就是指定主節(jié)點(diǎn)的world_size: 介紹都是說(shuō)是進(jìn)程, 實(shí)際就是機(jī)器的個(gè)數(shù), 例如兩臺(tái)機(jī)器一起訓(xùn)練的話, world_size就設(shè)置為2rank: 區(qū)分主節(jié)點(diǎn)和從節(jié)點(diǎn)的, 主節(jié)點(diǎn)為0, 剩余的為了1-(N-1), N為要使用的機(jī)器的數(shù)量, 也就是world_size

import os# 以下二選一, 第一個(gè)是使用gloo后端需要設(shè)置的, 第二個(gè)是使用nccl需要設(shè)置的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# 這一句就是和平時(shí)使用有點(diǎn)不一樣的地方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()