Capacitación multi-GPU y acumulación de gradientes en PyTorch como alternativa

¡Haz una fortuna con tu pequeña mano, dale un pulgar hacia arriba!

En este artículo [1] , primero entenderemos la diferencia entre los algoritmos Data Parallel (DP) y Distributed Data Parallel (DDP), luego explicaremos qué es Gradient Accumulation (GA), y finalmente mostraremos DDP y GA en las implementaciones de PyTorch y cómo conducen al mismo resultado.

Introducción

Cuando se entrena una red neuronal profunda (DNN), un hiperparámetro importante es el tamaño del lote. Por lo general, el tamaño del lote no debe ser demasiado grande, porque la red es propensa a sobreajustarse, pero no debe ser demasiado pequeño, porque conducirá a una convergencia lenta.

Cuando se trata de imágenes de alta resolución u otros tipos de datos que hacen un uso intensivo de la memoria, suponiendo que la mayor parte del entrenamiento de los modelos DNN grandes de hoy en día se realiza en GPU, el ajuste de lotes pequeños puede ser problemático dependiendo de la memoria de GPU disponible. Como dijimos, debido a que los lotes pequeños conducen a una convergencia lenta, existen tres métodos principales que podemos usar para aumentar el tamaño efectivo del lote:

  1. Ejecute modelos en paralelo en mini lotes utilizando varias GPU pequeñas: algoritmos DP o DDP
  2. Use una GPU más grande (caro)
  3. Acumule gradientes en varios pasos

Ahora veamos 1. y 3. con más detalle: si tiene la suerte de tener una GPU grande que pueda contener todos los datos que necesita, puede leer la sección DDP y ver cómo funciona en la sección de código completo implementado en PyTorch, saltándose así el resto.

Supongamos que queremos un tamaño de lote efectivo de 30, pero solo caben 10 puntos de datos en cada GPU (tamaño de minilote). Tenemos dos opciones: Datos Paralelos o Datos Distribuidos Paralelos:

Paralelismo de datos (DP)

Primero, definimos la GPU principal. Luego realizamos los siguientes pasos:

  1. Mueva 10 puntos de datos (minilotes) y una copia del modelo de la GPU principal a las otras 2 GPU
  2. Haga un pase hacia adelante en cada GPU y pase la salida a la GPU principal
  3. Calcule la pérdida total en la GPU principal, luego envíe la pérdida a cada GPU para calcular el gradiente de los parámetros
  4. 将梯度发送回Master GPU(这些是所有训练示例的梯度平均值),将它们相加得到整批30个的平均梯度
  5. 更新主 GPU 上的参数并将这些更新发送到其他 2 个 GPU 以进行下一次迭代

这个过程存在一些问题和低效率:

  • 数据-从主 GPU 传递,然后在其他 GPU 之间分配。此外,主 GPU 的利用率高于其他 GPU,因为总损失的计算和参数更新发生在主 GPU 上
  • 我们需要在每次迭代时同步其他 GPU 上的模型,这会减慢训练速度

分布式数据并行 (DDP)

引入分布式数据并行是为了改善数据并行算法的低效率。我们仍然采用与之前相同的设置 — 每批 30 个数据点,使用 3 个 GPU。差异如下:

  1. 它没有主 GPU
  2. 因为我们不再拥有主 GPU,所以我们直接从磁盘/RAM 以非重叠方式并行加载每个 GPU 上的数据 — DistributedSampler 为我们完成这项工作。在底层,它使用本地等级 (GPU id) 在 GPU 之间分配数据 - 给定 30 个数据点,第一个 GPU 将使用点 [0, 3, 6, ... , 27],第二个 GPU [1, 4, 7, .., 28] 和第三个 GPU [2, 5, 8, .. , 29]
n_gpu = 3
for i in range(n_gpu):
  print(np.arange(30)[i:30:n_gpu])
  1. 前向传递、损失计算和后向传递在每个 GPU 上独立执行,异步减少梯度计算平均值,然后在所有 GPU 上进行更新

由于DDP相对于DP的优点,目前优先使用DDP,因此我们只展示DDP的实现。

梯度累积

如果我们只有一个 GPU 但仍想使用更大的批量大小,另一种选择是累积一定数量的步骤的梯度,有效地累积一定数量的小批量的梯度,从而增加有效的批量大小。从上面的例子中,我们可以通过 3 次迭代累积 10 个数据点的梯度,以达到与我们在有效批量大小为 30 的 DDP 训练中描述的结果相同的结果。

DDP流程代码

下面我将仅介绍与 1 GPU 代码相比实现 DDP 时的差异。完整的代码可以在下面的一些部分找到。首先我们初始化进程组,允许不同进程之间进行通信。使用 int(os.environ[“LOCAL_RANK”]) 我们检索给定进程中使用的 GPU。

init_process_group(backend="nccl")
device = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(device)

然后,我们需要将模型包装在 DistributedDataParallel 中,以支持多 GPU 训练。

model = NeuralNetwork(args.data_size) 
model = model.to(device) 
 
if args.distributed:
  model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device])

最后一部分是定义我在 DDP 部分中提到的 DistributedSampler。

sampler = torch.utils.data.DistributedSampler(dataset)

培训的其余部分保持不变 - 我将在本文末尾包含完整的代码。

梯度累积代码

当反向传播发生时,在我们调用 loss.backward() 后,梯度将存储在各自的张量中。实际的更新发生在调用 optimizationr.step() 时,然后使用 optimizationr.zero_grad() 将张量中存储的梯度设置为零,以运行反向传播和参数更新的下一次迭代。

因此,为了累积梯度,我们调用 loss.backward() 来获取我们需要的梯度累积数量,而不将梯度设置为零,以便它们在多次迭代中累积,然后我们对它们进行平均以获得累积梯度迭代中的平均梯度(loss = loss/ACC_STEPS)。之后我们调用optimizer.step()并将梯度归零以开始下一次梯度累积。在代码中:

ACC_STEPS = dist.get_world_size() # == number of GPUs
# iterate through the data
for i, (idxs, row) in enumerate(loader):
  loss = model(row)  
  # scale loss according to accumulation steps
  loss = loss/ACC_STEPS
  loss.backward()
  # keep accumualting gradients for ACC_STEPS
  if ((i + 1) % ACC_STEPS == 0):
    optimizer.step()  
    optimizer.zero_grad()

代码

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
print(os.environ["CUDA_VISIBLE_DEVICES"])

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset, Sampler
import argparse
import torch.optim as optim 
import numpy as np
import random
import torch.backends.cudnn as cudnn
import torch.nn.functional as F

from torch.distributed import init_process_group
import torch.distributed as dist

class data_set(Dataset):
    
    def __init__(self, df):
        self.df = df
        
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, index):    
        
        sample = self.df[index]
        return index, sample
    
class NeuralNetwork(nn.Module):
    def __init__(self, dsize):
        super().__init__()
        self.linear =  nn.Linear(dsize, 1, bias=False)
        self.linear.weight.data[:] = 1.

    def forward(self, x):
        x = self.linear(x)
        loss = x.sum()
        return loss
    
            
class DummySampler(Sampler):
    def __init__(self, data, batch_size, n_gpus=2):
        self.num_samples = len(data)
        self.b_size = batch_size
        self.n_gpus = n_gpus

    def __iter__(self):
        ids = []
        for i in range(0, self.num_samples, self.b_size * self.n_gpus):
            ids.append(np.arange(self.num_samples)[i: i + self.b_size*self.n_gpus :self.n_gpus])
            ids.append(np.arange(self.num_samples)[i+1: (i+1) + self.b_size*self.n_gpus :self.n_gpus])
        return iter(np.concatenate(ids))

    def __len__(self):
        # print ('\tcalling Sampler:__len__')
        return self.num_samples

    

        
def main(args=None):
    
    d_size = args.data_size

    if args.distributed:
        init_process_group(backend="nccl")
        device = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(device)
    else:
        device = "cuda:0"

    # fix the seed for reproducibility
    seed = args.seed
        
        
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.benchmark = True
    
    # generate data
    data = torch.rand(d_size, d_size)
    
    model = NeuralNetwork(args.data_size)    
    model = model.to(device)  
    
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device])
        
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    dataset = data_set(data)

    if args.distributed:
        sampler = torch.utils.data.DistributedSampler(dataset, shuffle=False)
    else:
        # we define `DummySampler` for exact reproducibility with `DistributedSampler`
        # which splits the data as described in the article. 
        sampler = DummySampler(dataset, args.batch_size)
        
    loader = DataLoader(
                dataset,
                batch_size=args.batch_size,
                num_workers=0,
                pin_memory=True,
                sampler=sampler,
                shuffle=False,
                collate_fn=None,
            )          
    
    if not args.distributed:
        grads = []
    
    # ACC_STEPS same as GPU as we need to divide the loss by this number
    # to obtain the same gradient as from multiple GPUs that are 
    # averaged together
    ACC_STEPS = args.acc_steps 
    optimizer.zero_grad()
    
    for epoch in range(args.epochs):
        
        if args.distributed:
            loader.sampler.set_epoch(epoch)
            
        for i, (idxs, row) in enumerate(loader):

            if args.distributed:
                optimizer.zero_grad()
            
            row = row.to(device, non_blocking=True
            
            if args.distributed:
                rank = dist.get_rank() == 0
            else:
                rank = True
            
            loss = model(row)  
            
            if args.distributed:
                # does average gradients automatically thanks to model wrapper into 
                # `DistributedDataParallel`
                loss.backward()
            else:
                # scale loss according to accumulation steps
                loss = loss/ACC_STEPS
                loss.backward()
            
            if i == 0 and rank:
                print(f"Epoch {epoch} { 100 * '='}")

            if not args.distributed:
                if (i + 1) % ACC_STEPS == 0# only step when we have done ACC_STEPS
                    # acumulate grads for entire epoch
                    optimizer.step()  
                    optimizer.zero_grad()
            else:
                optimizer.step() 

                
        if not args.distributed and args.verbose:
            print(100 * "=")
            print("Model weights : ", model.linear.weight)
            print(100 * "=")
        elif args.distributed and args.verbose and rank:
            print(100 * "=")
            print("Model weights : ", model.module.linear.weight)
            print(100 * "=")

    


if __name__ == "__main__":
    
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--distributed', action='store_true',)
    parser.add_argument('--seed', default=0, type=int) 
    parser.add_argument('--epochs', default=2, type=int) 
    parser.add_argument('--batch_size', default=4, type=int) 
    parser.add_argument('--data_size', default=16, type=int) 
    parser.add_argument('--acc_steps', default=3, type=int) 
    parser.add_argument('--verbose', action='store_true',)
    
    args = parser.parse_args()
    
    print(args)

    main(args)

总结

在本文中,我们简要介绍并直观地介绍了 DP、DDP 算法和梯度累积,并展示了如何在没有多个 GPU 的情况下增加有效批量大小。需要注意的一件重要事情是,即使我们获得相同的最终结果,使用多个 GPU 进行训练也比使用梯度累积要快得多,因此如果训练速度很重要,那么使用多个 GPU 是加速训练的唯一方法。

Reference

[1]

Source: https://towardsdatascience.com/multiple-gpu-training-in-pytorch-and-gradient-accumulation-as-an-alternative-to-it-e578b3fc5b91

本文由 mdnice 多平台发布

Supongo que te gusta

Origin blog.csdn.net/swindler_ice/article/details/132029502
Recomendado
Clasificación