Habilidades de aceleración de entrenamiento de PyTorch

Dado que los programas recientes tienen requisitos relativamente altos de velocidad y quieren producir resultados rápidamente, aprendí especialmente operaciones aritméticas y paralelas de precisión mixta. Dado que ha habido muchos artículos que presentan principios relacionados, este artículo solo describe cómo aplicar PyTorch para lograr cálculos de precisión híbridos. , el paralelismo de datos y los cálculos distribuidos no introducen específicamente los principios.

Precisión mixta

El entrenamiento automático de precisión mixta (auto Mixed Precision, AMP) puede reducir en gran medida el costo del entrenamiento y aumentar la velocidad del entrenamiento. Antes de esto, los cálculos automáticos de precisión mixta utilizaban herramientas Apex desarrolladas por NVIDIA. A partir de PyTorch 1.6.0, PyTorch ya ha enviado su propio módulo AMP, por lo que lo siguiente presentará principalmente el uso del módulo de amplificador propio de PyTorch .

## 导入amp工具包 
from torch.cuda.amp import autocast, GradScaler

model.train()

## 对梯度进行scale来加快模型收敛，
## 因为float16梯度容易出现underflow（梯度过小）
scaler = GradScaler()

batch_size = train_loader.batch_size
num_batches = len(train_loader)
end = time.time()
for i, (images, target) in tqdm.tqdm(
    enumerate(train_loader), ascii=True, total=len(train_loader)
):
    # measure data loading time
    data_time.update(time.time() - end)
    optimizer.zero_grad()
    if args.gpu is not None:
        images = images.cuda(args.gpu, non_blocking=True)

    target = target.cuda(args.gpu, non_blocking=True)
    # 自动为GPU op选择精度来提升训练性能而不降低模型准确度
    with autocast():
    # compute output
        output = model(images)

        loss = criterion(output, target)

    scaler.scale(loss).backward()
    # optimizer.step()
    scaler.step(optimizer)
    scaler.update()

Paralelo de datos

Cuando el servidor tiene una sola máquina con múltiples tarjetas, para lograr la aceleración del modelo (tal vez porque una GPU no es suficiente), se puede usar una sola máquina con múltiples tarjetas para entrenar el modelo. Para lograr este objetivo, debemos encontrar una manera de permitir que un modelo se distribuya en múltiples GPU para entrenamiento.

En PyTorch , nn.DataParallel nos proporciona una interfaz simple, que puede realizar fácilmente la paralelización del modelo. Solo necesitamos usar nn.DataParallel para empaquetar el modelo. Después de configurar algunos parámetros, el modelo se puede implementar fácilmente. tarjeta paralela.

# multigpu表示显卡的号码
multigpu = [0,1,2,3,4,5,6,7] 
# 设置主GPU,用来汇总模型的损失函数并且求导，对梯度进行更新
torch.cuda.set_device(args.multigpu[0])
# 模型的梯度全部汇总到gpu[0]上来
model = torch.nn.DataParallel(model, device_ids=args.multigpu).cuda(
		args.multigpu[0]
      	)

nn.DataParallel utiliza aritmética de precisión mixta

    nn.DataParallel requiere alguna configuración especial para realizar operaciones de precisión mixta en el modelo; de lo contrario, el modelo no puede lograr la paralelización de datos.
    Autocast está diseñado para ser "subproceso local", por lo que solo configurar el área de autocast en el subproceso principal no funciona. Tomado prestado de la referencia
    aquí para dar la operación incorrecta:

model = MyModel() 
dp_model = nn.DataParallel(model)

with autocast():     # dp_model's internal threads won't autocast.
     #The main thread's autocast state has no effect.     
     output = dp_model(input)     # loss_fn still autocasts, but it's too late...
     loss = loss_fn(output)

Hay dos soluciones, que se presentan a continuación:
1. Agregue una función decorativa a la función de avance del módulo de modelo.

MyModel(nn.Module):
    ...
    @autocast()
    def forward(self, input):
       ...

2. Otra postura correcta es establecer el área de transmisión automática hacia adelante:

MyModel(nn.Module):
   ...
   def forward(self, input):
       with autocast():
           ...

Después de operar la función de avance, use autocast en el hilo principal

model = MyModel()
dp_model = nn.DataParallel(model)

with autocast():
   output = dp_model(input)
   loss = loss_fn(output)

nn.DataParallel desventajas

En cada lote de entrenamiento, el módulo nn.DataParallel retransmitirá todas las pérdidas a gpu [0]. La transmisión de datos de varios Gs y los cálculos de pérdidas deben completarse en una tarjeta gráfica, lo cual es muy fácil La carga de la tarjeta gráfica es desigual y, a menudo, se puede ver que la carga de gpu [0] es significativamente mayor que la de otras gpus. Además, la velocidad de transmisión de datos de la tarjeta gráfica provocará un gran cuello de botella en la velocidad de entrenamiento del modelo, lo que obviamente no es razonable.
Entonces, a continuación, presentaremos, el principio específico puede referirse a la operación de múltiples tarjetas en una sola máquina (DataParallel distribuido, precisión mixta, Horovod)

Computación distribuída

nn.DistributedDataParallel : control multiproceso multi-GPU, entrenar el modelo juntos.

ventaja

Cada proceso controla una GPU, que puede garantizar que el cálculo del modelo no se vea afectado por la comunicación entre las tarjetas gráficas y puede hacer que la carga de cada tarjeta gráfica sea relativamente uniforme. Pero en comparación con una sola tarjeta o una sola tarjeta ( nn.DataParallel ), existen varios problemas

1. Sincronice los parámetros del modelo en diferentes GPU, especialmente BatchNormalization
2. Indique a cada proceso su propia ubicación, qué GPU usar y especifique con el parámetro args.local_rank
. 3. Al obtener datos, cada proceso debe asegurarse de que lo que obtiene son datos diferentes (DistributedSampler)

Introducción al uso

Procedimiento de inicio
Dado que el blogger solo ha practicado la operación de múltiples tarjetas en una sola máquina en la actualidad, presenta principalmente la operación de múltiples tarjetas en una sola máquina. A diferencia de la ejecución simple habitual de programas de Python, necesitamos usar el lanzador torch.distributed.launch que viene con PyTorch para iniciar el programa.

# 其中CUDA_VISIBLE_DEVICES指定机器上显卡的数量
# nproc_per_node程序进程的数量
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

Configurar el programa principal

parser.add_argument('--local_rank', type=int, default=0，help='node rank for distributed training')
# 配置local_rank参数，告诉每个进程自己的位置，要使用哪张GPU

Inicializar la forma de comunicación de la tarjeta gráfica y adquisición de parámetros.

# 为这个进程指定GPU
torch.cuda.set_device(args.local_rank)
# 初始化GPU通信方式NCLL和参数的获取方式，其中env表示环境变量
# PyTorch实现分布式运算是通过NCCL进行显卡通信的
torch.distributed.init_process_group(
    backend='nccl',
    rank=args.local_rank
)

Reconfigurar DataLoader


kwargs = {
    
    "num_workers": args.workers, "pin_memory": True} if use_cuda else {
    
    }

train_sampler = DistributedSampler(train_dataset)
self.train_loader = torch.utils.data.DataLoader(
            train_dataset, 
            batch_size=args.batch_size, 
            sampler=train_sampler,  
            **kwargs
        )

# 注意，由于使用了Sampler方法，dataloader中就不能加shuffle、drop_last等参数了
'''
PyTorch dataloader.py 192-197 代码
        if batch_sampler is not None:
            # auto_collation with custom batch_sampler
            if batch_size != 1 or shuffle or sampler is not None or drop_last:
                raise ValueError('batch_sampler option is mutually exclusive '
                                 'with batch_size, shuffle, sampler, and '
                                 'drop_last')'''

pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。

Inicialización del modelo

torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)
model.to(device)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = torch.nn.parallel.DistributedDataParallel(
        model,
        device_ids=[args.local_rank],
        output_device=args.local_rank,
        find_unused_parameters=True,
        )
torch.backends.cudnn.benchmark=True 
# 将会让程序在开始时花费一点额外时间，为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速
# DistributedDataParallel可以将不同GPU上求得的梯度进行汇总，实现对模型GPU的更新

DistributedDataParallel puede resumir los gradientes obtenidos en diferentes GPU para actualizar el modelo de GPU

Sincronizar la capa BatchNormalization

Para las tareas de entrenamiento que consumen más memoria de video, el tamaño de lote relativo en una sola tarjeta es a menudo demasiado pequeño, lo que afecta el efecto de convergencia del modelo. Sincronización entre tarjetas La normalización por lotes puede utilizar muestras globales para la normalización, lo que equivale a "aumentar" el tamaño del lote, de modo que el efecto de entrenamiento ya no se vea afectado por la cantidad de GPU utilizadas. Consulte la operación autónoma de varias tarjetas (DataParallel distribuido, precisión mixta, Horovod).
Afortunadamente, en la versión reciente de Pytorch, PyTorch ha comenzado a admitir de forma nativa la sincronización de la capa BatchNormalization.

antorcha.nn.SyncBatchNorm
torch.nn.SyncBatchNorm.convert_sync_batchnorm : convierte automáticamente la capa de BatchNorm-alization en torch.nn.SyncBatchNorm para lograr la sincronización de la capa BatchNormalization en diferentes GPU

Para una implementación específica, consulte el código de inicialización del modelo.

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

Semilla aleatoria para la inicialización del modelo de sincronización

En la actualidad, no he probado la situación de usar diferentes semillas aleatorias en diferentes procesos. Para estar seguro, se recomienda asegurarse de que la semilla aleatoria para cada inicialización de modelo sea la misma para garantizar que los modelos en cada proceso de GPU estén sincronizados.

para resumir

De pie sobre los hombros de gigantes aceleré hace tiempo el modelo de autoaprendizaje, pisé muchos boxes y finalmente agregué el desfile, finalmente resumí algunos códigos específicos y me referí a muchos otros blogs. Espero ser de alguna ayuda para todos.

Cita (sin importar el anverso y el reverso)