PyTorch Distributed Tutorials(6) PyTorch 1.0 Distributed Trainer with Amazon AWS

---

0. 前言

• 官方链接
• 目标：在两台AWS上运行分布式训练代码。

---

1. AWS 配置

• 新建实例。
• 重点设置安全组。
• 记录实例的内部ip与外部ip。

---

2. 环境配置

• 每个实例需要分别单独处理。
• 基本工作：
  • 建立conda环境，安装pytorch与vision。
  • 重要：设置NCCL socket的infterface name。
    • 通过 ifconfig 查看当前实例的interface name，即内部ip对应的name。
    • 然后通过环境变量 NCCL_SOCKET_IFNAME 进行设置。假设name为ens3，则需要设置 export NCCL_SOCKET_IFNAME=ens3
    • 所有节点都要这么处理。
• 由于没有设置分布式文件系统（shared filesystem），所以每个实例都需要有各自的数据集以及代码的副本。

---

3. 分布式训练代码

• 本文以 pytorch imagenet example 为例进行说明，主要代码都是来自这里。

3.1. imports & helper function

• 其实就是将所有import语句以及计算metrics中的auucary以及平均值的函数/类。

import time
import sys
import torch

import torch.nn as nn
import torch.nn.parallel
import torch.distributed as dist
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models

from torch.multiprocessing import Pool, Process

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res

3.2. 训练代码

• 包括两个函数（看着代码本身没什么不一样的）
  • train 用于训练一个epoch，可能就是一个 .cuda(non_blocking=True)
    • 这个参数设置为True后，allows asynchronous GPU copies of the data, meaning transfers can be overlapped with computation，允许异步使用GPU数据副本。此时数据传输与计算可以同时进行。
  • adjust_learning_rate 用于更新训练过程中的学习率。

def train(train_loader, model, criterion, optimizer, epoch):

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    for i, (input, target) in enumerate(train_loader):

        # measure data loading time
        data_time.update(time.time() - end)

        # Create non_blocking tensors for distributed training
        input = input.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)

        # compute output
        output = model(input)
        loss = criterion(output, target)

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), input.size(0))
        top1.update(prec1[0], input.size(0))
        top5.update(prec5[0], input.size(0))

        # compute gradients in a backward pass
        optimizer.zero_grad()
        loss.backward()

        # Call step of optimizer to update model params
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % 10 == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   epoch, i, len(train_loader), batch_time=batch_time,
                   data_time=data_time, loss=losses, top1=top1, top5=top5))

def adjust_learning_rate(initial_lr, optimizer, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    lr = initial_lr * (0.1 ** (epoch // 30))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

3.3. 验证代码

• 好像没有什么分布式相关的（除了non_blocking）

def validate(val_loader, model, criterion):

    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        end = time.time()
        for i, (input, target) in enumerate(val_loader):

            input = input.cuda(non_blocking=True)
            target = target.cuda(non_blocking=True)

            # compute output
            output = model(input)
            loss = criterion(output, target)

            # measure accuracy and record loss
            prec1, prec5 = accuracy(output, target, topk=(1, 5))
            losses.update(loss.item(), input.size(0))
            top1.update(prec1[0], input.size(0))
            top5.update(prec5[0], input.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % 100 == 0:
                print('Test: [{0}/{1}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                       i, len(val_loader), batch_time=batch_time, loss=losses,
                       top1=top1, top5=top5))

        print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
              .format(top1=top1, top5=top5))

    return top1.avg

3.4. 设置输入数据相关参数

print("Collect Inputs...")

# Batch Size for training and testing
batch_size = 32

# 每个进程process中dataloaders的worker process数量
# Number of additional worker processes for dataloading
workers = 2

# Number of epochs to train for
num_epochs = 2

# Starting Learning Rate
starting_lr = 0.1

# 分布式训练节点数量
# Number of distributed processes
world_size = 4

# Distributed backend type
dist_backend = 'nccl'

# 设置 process group 的 initialization 方法所在ip以及port
# 也可能是分布式文件系统中一个不存在的文件路径
# 一般就是内部ip加上node0使用的port
# Url used to setup distributed training
dist_url = "tcp://172.31.22.234:23456"

3.5. 初始化 process group

• 初始化 Process group 是分布式训练中非常重要的一个步骤，这是初始化 torch.distributed 的第一步。
• 初始化方法是 torch.distributed.init_process_group(backend, init_method, rank, world_size)
  • backend 是 NCCL/Gloo/MPI 三选一，一般GPU也就是用NCCL。
  • init_method 可以是 ip+port，也可以是分布式文件系统中一个不存在的文件。
    • 注意，不管是ip+port或分布式文件系统中的文件，都要求所有节点都能访问。
  • rank 用来指定当前节点的rank，如果没有指定参数则可以通过环境变量 RANK 指定。
  • world_size GPU数量，如果没有指定参数，则可以通过环境变量WORLD_SIZE指定。
• 另外非常重要的一步就是设置 local rank
  • 比如有两个节点，每个节点8个GPU，则work size就是16，且每个节点中的local rank都是0-7。

print("Initialize Process Group...")
# Initialize Process Group
# v1 - init with url
dist.init_process_group(backend=dist_backend, init_method=dist_url, rank=int(sys.argv[1]), world_size=world_size)
# v2 - init with file
# dist.init_process_group(backend="nccl", init_method="file:///home/ubuntu/pt-distributed-tutorial/trainfile", rank=int(sys.argv[1]), world_size=world_size)
# v3 - init with environment variables
# dist.init_process_group(backend="nccl", init_method="env://", rank=int(sys.argv[1]), world_size=world_size)


# Establish Local Rank and set device on this node
local_rank = int(sys.argv[2])
dp_device_ids = [local_rank]
torch.cuda.set_device(local_rank)

3.6. 模型初始化

• 基本步骤：
  • 建立普通的单GPU模型，即初始化模型，然后放入GPU中。
  • 之后，使用 DistributedDataParallel 包装原有模型，设置 local rank。
    • 注意，该类不仅实现了数据分布式处理，还处理了梯度的分布式处理。
  • 定义了损失函数与optimizer。

print("Initialize Model...")
# Construct Model
model = models.resnet18(pretrained=False).cuda()
# Make model DistributedDataParallel
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=dp_device_ids, output_device=local_rank)

# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), starting_lr, momentum=0.9, weight_decay=1e-4)

3.7. 初始化 dataloaders

• 基本步骤
  • 定义普通的数据集对象。
  • 定义 distributedsampler 从而实现分布式。
  • 构建dataloader，将上面的distributed sampler作为构造函数的一部分。

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 present
trainset = datasets.STL10(root='./data', split='train', download=True, transform=transform)
valset = datasets.STL10(root='./data', split='test', download=True, transform=transform)

# 只有在 init_process_group 执行完成后才能执行这一步
# 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 steps
train_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)

3.8. 训练loop代码

• 实现若干epoch的训练，没什么花头

best_prec1 = 0

for epoch in range(num_epochs):
    # Set epoch count for DistributedSampler
    train_sampler.set_epoch(epoch)

    # Adjust learning rate according to schedule
    adjust_learning_rate(starting_lr, optimizer, epoch)

    # train for one epoch
    print("\nBegin Training Epoch {}".format(epoch+1))
    train(train_loader, model, criterion, optimizer, epoch)

    # evaluate on validation set
    print("Begin Validation @ Epoch {}".format(epoch+1))
    prec1 = validate(val_loader, model, criterion)

    # remember best prec@1 and save checkpoint if desired
    # is_best = prec1 > best_prec1
    best_prec1 = max(prec1, best_prec1)

    print("Epoch Summary: ")
    print("\tEpoch Accuracy: {}".format(prec1))
    print("\tBest Accuracy: {}".format(best_prec1))

---

4. 开始分布式训练

• 环境：两个节点，每个节点2个GPU。

  • 所以 world_size 为4。
  • rank可以理解为GPU数量，local rank为当前节点中gpu编号。

• 基本流程（第一个参数为 rank，第二个参数为 local_rank）：

  • 在第一个节点中运行：python main.py 0 0
  • 在第二个节点中运行：python main.py 1 1
  • 在第三个节点中运行：python main.py 2 0
  • 在第四个节点中运行：python main.py 3 1

• 可以通过 launcher 工具或 torch.multiprocessing.spawn 实现类似上面四个语句的功能。