[人工智能]PyTorch DistributedDataParallel分布式训练模板（可以执行运行，查看效果）

引言

• PyTorch多GPU训练有两种方式：
  • DataParallel: 代码简单，训练太慢
  • DistributedDataParallel: 代码较复杂，训练速度很快
• 今天主要提供DistributedDataParallel的模板，可以直接使用

代码

• 首先，验证一下自己机器是否可以使用DistributedDataParallel

  import torch
  
  torch.distributed.is_available()
  # True → 可以	
  

• 模板代码(可以直接运行的)

  import torch
  import argparse
  import torch.distributed as dist
  import torch.nn as nn
  import torch.nn.functional as F
  from torchvision import datasets, transforms
  import torch.optim as optim
  
  
  class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
          self.conv1 = nn.Conv2d(1, 32, 3, 1)
          self.conv2 = nn.Conv2d(32, 64, 3, 1)
          self.dropout1 = nn.Dropout(0.25)
          self.dropout2 = nn.Dropout(0.5)
          self.fc1 = nn.Linear(9216, 128)
          self.fc2 = nn.Linear(128, 10)
  
      def forward(self, x):
          x = self.conv1(x)
          x = F.relu(x)
          x = self.conv2(x)
          x = F.relu(x)
          x = F.max_pool2d(x, 2)
          x = self.dropout1(x)
          x = torch.flatten(x, 1)
          x = self.fc1(x)
          x = F.relu(x)
          x = self.dropout2(x)
          x = self.fc2(x)
          output = F.log_softmax(x, dim=1)
          return output
  
  torch.manual_seed(0)
  torch.backends.cudnn.deterministic = True
  torch.backends.cudnn.benchmark = False
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--local_rank', default=0, type=int,
                      help='node rank for distributed training')
  args = parser.parse_args()
  
  torch.cuda.set_device(args.local_rank)
  # world_size 等于使用gpu的个数
  dist.init_process_group(backend='nccl', init_method='env://',
                          world_size=3, rank=args.local_rank)
  
  transform=transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize((0.1307,), (0.3081,))
      ])
  
  train_dataset = datasets.MNIST('../dataset', train=True, download=True,
                      transform=transform)
  val_dataset = datasets.MNIST('../dataset', train=False,
                      transform=transform)
  train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  train_loader = torch.utils.data.DataLoader(train_dataset,
                                             batch_size=4,
                                             sampler=train_sampler)
  
  device = torch.device('cuda')
  model = Net().to(device)
  model = torch.nn.parallel.DistributedDataParallel(model,
                                                    device_ids=[args.local_rank],
                                                    	output_device=args.local_rank)
  
  optimizer = optim.SGD(model.parameters(), lr=1e-3)
  criterion = nn.CrossEntropyLoss()
  for epoch in range(100):
     for batch_idx, (data, target) in enumerate(train_loader):
          images = data.cuda(non_blocking=True)
          target = target.cuda(non_blocking=True)
  
          output = model(images)
          loss = criterion(output, target)
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          print(loss.item())
  

• 运行

  export MASTER_ADDR=localhost
  export MASTER_PORT=5678
  CUDA_VISIBLE_DEVICES=1,2,3 python -m torch.distributed.launch --nproc_per_node=3 demo.py
  

参考资料

• pytorch-examples-mnist
• pytorch-distributed