Precisión mixta automática utilizando PyTorch

Introducción

Los modelos de aprendizaje profundo más grandes requieren más potencia de cómputo y recursos de memoria. Se ha logrado un entrenamiento más rápido de redes neuronales profundas mediante el desarrollo de nuevas técnicas. En lugar de FP32 (formato de números flotantes de precisión completa), se puede utilizar FP16 (formato de números flotantes de mitad de precisión), y los investigadores han descubierto que su uso en conjunto es una opción mejor.

La precisión mixta permite el entrenamiento con mitad de precisión mientras conserva gran parte de la precisión de la red de precisión simple. El término “técnica de precisión mixta” se refiere a que este método hace uso de ambas representaciones de precisión simple y mitad.

En este resumen del entrenamiento de Precisión Mixta Automática (AMP) con PyTorch, demostramos cómo funciona la técnica, caminando paso a paso por el proceso de uso de AMP, y discutimos aplicaciones más avanzadas de técnicas de AMP con estructuras de código para que los usuarios integruen posteriormente en su propio código.

Requisitos previos

Conocimientos básicos de PyTorch: Familiarización con PyTorch, incluyendo sus conceptos básicos como tensores, módulos y el bucle de entrenamiento.

Comprensión de los fundamentos de la aprendizaje profundo: Conceptos como redes neuronales, propagación hacia atrás y optimización.

Conocimiento de Entrenamiento con Precisión Mixta: Conocimiento de los beneficios y desventajas del entrenamiento con precisión mixta, incluyendo el uso reducido de memoria y una computación más rápida.

Acceso a Hardware Compatible: Un GPU que soporte precisión mixta, como los GPU de NVIDIA con Cores Tensor (por ejemplo, arquitecturas Volta, Turing, Ampere).

Configuración de Python y CUDA: Un entorno de Python funcional con PyTorch instalado y CUDA configurado para la aceleración GPU.

Resumen de la Precisión Mixta

Al igual que la mayoría de los frameworks de aprendizaje profundo, PyTorch normalmente se entrena con datos de punto flotante de 32 bits (FP32). Sin embargo, los FP32 no son siempre necesarios para el éxito. Es posible utilizar un punto flotante de 16 bits para algunas operaciones, donde FP32 consume más tiempo y memoria.

En consecuencia, los ingenieros de NVIDIA desarrollaron una técnica que permite realizar el entrenamiento con precisión mixta en FP32 para un pequeño número de operaciones, mientras que la mayoría de la red se ejecuta en FP16.

• Convierta el modelo para utilizar el tipo de dato float16 donde sea posible.
• Mantenga los pesos master float32 para acumular actualizaciones de peso cada iteración.
• El uso de escalado de pérdida para preservar valores de gradientes pequeños.

Entrenamiento de precisión mixta en PyTorch

PyTorch ofrece una gran cantidad de características integradas para el entrenamiento de precisión mixta.
Los parámetros de un módulo se convierten a FP16 cuando se llama al método .half(), y los datos de un tensor se convierten a FP16 cuando se llama a .half(). Se utilizará la aritmética rápida FP16 para ejecutar cualquier operación en estos módulos o tensores. Las bibliotecas de matemáticas de NVIDIA (cuBLAS y cuDNN) están bien soportadas por PyTorch. Los datos de la tubería FP16 se procesan utilizando las Cores Tensor para realizar GEMMs y convulsiones. Para emplear Cores Tensor en cuBLAS, las dimensiones de un GEMM ([M, K] x [K, N] -> [M, N]) deben ser múltiplos de 8.

Introducción a Apex

Las utilidades de precisión mixta de Apex están destinadas a aumentar la velocidad de entrenamiento mientras mantiene la precisión y estabilidad del entrenamiento de single-precision. Apex puede realizar operaciones en FP16 o FP32, manejar automáticamente la conversión de parámetros maestros y escalar automáticamente las pérdidas.

Apex fue creado para facilitar que investigadores incluyan en sus modelos el entrenamiento de precisión mixta. Amp, que es una abreviatura de Automatic Mixed-Precision, es una de las características de Apex, una extensión de PyTorch ligera. Solo son unas líneas más en sus redes las que los usuarios necesitan para aprovechar el entrenamiento de precisión mixta con Amp. Apex fue lanzado en CVPR 2018, y cabe destacar que la comunidad de PyTorch ha mostrado un fuerte apoyo a Apex desde su lanzamiento.

Mediante sólo pequeños cambios en el modelo en ejecución, Amp hace que no tengas que preocuparte por los tipos mixtos mientras creas o ejecutas tu script. Las suposiciones de Amp pueden no encajar tan bien en modelos que utilizan PyTorch de maneras inusuales, pero hay ganchos para ajustar esas suposiciones según sea necesario.

Amp ofrece todos los beneficios del entrenamiento de precisión mixta sin la necesidad de escalado de pérdidas o conversiones de tipos que deban ser manejadas explícitamente. El sitio web de GitHub de Apex contiene instrucciones para el proceso de instalación, y su documentación de API oficial puede encontrarse aquí.

Cómo funcionan los Amps

Amp utiliza un paradigma de lista blanca/lista negra a nivel lógico. Las operaciones de tensor en PyTorch incluyen funciones de redes neuronales como torch.nn.functional.conv2d, funciones matemáticas sencillas como torch.log, y métodos de tensor como torch.Tensor. add__. Existen tres categorías principales de funciones en este universo:

• Lista blanca: Funciones que podrían beneficiarse de la velocidad del cálculo en FP16. Aplicaciones típicas incluyen la multiplicación de matrices y la convulsión.
• Lista negra: Los inputs deben estar en FP32 para funciones donde los 16 bits de precisión pueden no ser suficientes.
• Todo lo demás (cualquieras funciones que queden): Funciones que pueden ejecutarse en FP16, pero la expensa de una conversión de FP32 a FP16 para ejecutarlas en FP16 no es rentable ya que la aceleración no es significativa.

La tarea de Amp es sencilla, al menos en teoría. Amp determina si una función de PyTorch está en la lista blanca, en la lista negra o en ninguna, antes de llamarla. Todos los argumentos deben convertirse a FP16 si está en la lista blanca o a FP32 si está en la lista negra. Si no está en ninguna, solo asegúrese de que todos los argumentos son del mismo tipo. Esta política no es tan fácil de aplicar en la realidad como parece.

Uso de Amp en Conjunción con un Modelo de PyTorch

Para incluir Amp en un script de PyTorch actual, siga los siguientes pasos:

• Use la biblioteca Apex para importar Amp.
• Inicialice Amp para que pueda realizar los cambios requeridos al modelo, al optimizador y a las funciones internas de PyTorch.
• Tenga en cuenta dónde se produce la backpropagación (.backward()) para que Amp pueda escalar simultáneamente la pérdida y eliminar el estado por iteración.

Paso 1

Existe solamente una línea de código para el primer paso:

from apex import amp

Paso 2

El modelo de red neuronal y el optimizador utilizados para el entrenamiento deben ya estar especificados para completar este paso, que es una sola línea de longitud.

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

Las configuraciones adicionales le permiten ajustar con precisión las escalas de los tensores y tipos de operaciones del Amp. La función amp.initialize() acepta muchos parámetros, entre los cuales especificaremos solo tres de ellos:

• models (torch.nn.Module o lista de torch.nn.Modules) – Modelos a modificar/castear.
• optimizadores (opcional, torch.optim.Optimizer o lista de torch.optim.Optimizers) – Optimizadores para modificar/castear. OBLIGATORIO para el entrenamiento, opcional para la inferencia.
• opt_level (str, opcional, predeterminado=“O1”) – Nivel de optimización de precisión pura o mezclada. Los valores aceptados son “O0”, “O1”, “O2” y “O3”, explicados en detalle arriba. Hay cuatro niveles de optimización:

O0 para el entrenamiento en FP32: Esto no realiza ninguna operación. No hay que preocuparse por esto ya que su modelo de entrada debería ser ya en FP32, y O0 puede ayudar a establecer un punto de referencia para la precisión.

O1 para Precisión Mezclada (recomendado para el uso típico): Modificar todos los métodos de Tensor y Torch para utilizar un esquema de casting de entrada lista-blanco/negro. En FP16, se realizan operaciones en la lista blanca, como por ejemplo las operaciones amigas de Tensor Core como GEMMs y convoluciones. Softmax, por ejemplo, es una operación en la lista negra que requiere precisión FP32. A menos que se establezca de otra manera, O1 también emplea escalado dinámico de pérdidas.

O2 para Precisión Mezclada “Casi FP16”: O2 convierte los pesos del modelo en FP16, patcha el método forward del modelo para convertir los datos de entrada a FP16, mantiene las batchnorms en FP32, mantiene los pesos maestros en FP32, actualiza los grupos de parámetros del optimizador para que el optimizador.step() actúe directamente sobre los pesos en FP32 y implementa escalado dinámico de pérdidas (a menos que se sobreescriba). A diferencia de O1, O2 no patcha funciones de Torch o métodos de Tensor.

O3 para entrenamiento con FP16: O3 puede no ser tan estable como O1 y O2 en cuanto a precisión mixta real. Consecuentemente, podría ser ventajoso establecer una velocidad de referencia para su modelo, contra la cual la eficiencia de O1 y O2 pueda evaluarse.
La sobreescritura de propiedad adicional keep_batchnorm_fp32=True en O3 podría ayudarle a determinar la “velocidad de la luz” si su modelo emplea normalización por lotes, habilitando la normalización por lotes de cudnn.

O0 y O3 no son precisión mixta real, pero ayudan a establecer las líneas de base de precisión y velocidad, respectivamente. Una implementación de precisión mixta se define como O1 y O2.
Puede probar ambas y ver cuál mejora más el rendimiento y la precisión para su modelo particular.

Paso 3

Asegúrese de identificar dónde ocurre la pasada hacia atrás en su código.
Aparecerán unas líneas de código que se parecen a esto:

loss = criterion(…)
loss.backward()
optimizer.step()

Paso 4

Usando el gestor de contexto Amp, puede habilitar la escalado de pérdida simplemente enrollando la pasada hacia atrás:

loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

Eso es todo. Ahora puede volver a ejecutar su script con el entrenamiento de precisión mixta habilitado.

Captura de llamadas de función

PyTorch carece de un objeto de modelo estático o gráfico para afianzarse y insertar las conversiones mencionadas anteriormente, ya que es tan flexible y dinámico. Mediante el “monkey patching” de las funciones requeridas, Amp puede interceptar y convertir parametros dinámicamente.

Por ejemplo, puede usar el código de abajo para asegurarse que los argumentos del método torch.nn.functional.linear siempre se convierten en fp16:

orig_linear = torch.nn.functional.linear
def wrapped_linear(*args):
 casted_args = []
  for arg in args:
    if torch.is_tensor(arg) and torch.is_floating_point(arg):
      casted_args.append(torch.cast(arg, torch.float16))
    else:
      casted_args.append(arg)
  return orig_linear(*casted_args)
torch.nn.functional.linear = wrapped_linear

Aunque Amp pueda agregar refinerías para hacer que el código sea más resistente, llamar a Amp.init() efectivamente causa que los monkeys patches se inserjan en todas las funciones relevantes de PyTorch para que los argumentos se conviertan correctamente en tiempo de ejecución.

Minimización de Conversiones

Cada peso solo se convierte de FP32 a FP16 una sola vez por iteración, ya que Amp mantiene un caché interno de todas las conversiones de parámetros y las reutiliza según sea necesario. En cada iteración, el gestor de contexto para el paso en reversa indica a Amp cuándo limpiar el caché.

Autocasting y Escalado de Gradientes Utilizando PyTorch

“Entrenamiento de precisión mixta automatizado” se refiere a la combinación de torch.cuda.amp.autocast y torch.cuda.amp.GradScaler. Mediante torch.cuda.amp.autocast, podría configurar el autocast solo para ciertas áreas. El autocast selecciona automáticamente la precisión para las operaciones del GPU para optimizar la eficiencia mientras mantiene la precisión.

Las instancias de torch.cuda.amp.GradScaler hacen que sea más fácil realizar los pasos de escalado de gradientes. El escalado de gradientes reduce el underflow de gradientes, lo que ayuda a las redes con gradientes de float16 alcanzar una mejor convergencia.

Aquí hay un código para demostrar cómo utilizar autocast() para obtener una precisión mixta automatizada en PyTorch:

# Crea modelo y optimizador en precisión predeterminada
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Crea un escalador de gradientes una sola vez al inicio del entrenamiento.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Ejecuta el forward pass con autocast.
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # Las operaciones de retroceso se ejecutan en el mismo tipo de dato que autocast eligió para las operaciones correspondientes de adelante.
        scaler.scale(loss).backward()

        # scaler.step() primero desescala los gradientes de los parámetros asignados del optimizador.
   
        scaler.step(optimizer)

        # Actualiza la escala para la próxima iteración.
        scaler.update()

Si la pasada frontal para una operación específica tiene entradas de float16, entonces la pasada inversa para esta operación produce gradientes de float16, y float16 puede no ser capaz de expresar gradientes con magnitudes pequeñas.

Se perderán las actualizaciones de los parámetros relacionados si estos valores se descargan a cero («underflow»).

La escalada de gradientes es una técnica que utiliza un factor de escala para multiplicar las pérdidas de la red y luego realiza una pasada inversa sobre la pérdida escalada para evitar el underflow. También es necesario escalar los gradientes que fluyen hacia atrás a través de la red por este mismo factor. Consecuentemente, los valores de los gradientes tienen una magnitud mayor, lo que evita que se descargan a cero.

Antes de actualizar los parámetros, cada gradiente de parámetro (atributo .grad) debería desescalarse para que el factor de escala no interfiera con la tasa de aprendizaje. Tanto autocast como GradScaler se pueden usar de forma independiente ya que son modulares.

Trabajar con Gradientes No Escalados

Recortar Gradientes

Podemos escalar todos los gradientes utilizando el método Scaler.scale(Loss).backward(). Las propiedades .grad de los parámetros entre backward() y scaler.step(optimizer) deben ser desescaladas antes de cambiarlas o inspeccionarlas. Si desea limitar la norma global (ver torch.nn.utils.clip_grad_norm_()) o la magnitud máxima (ver torch.nn.utils.clip_grad_value_()) del conjunto de gradientes a ser menor o igual a un cierto valor (una determinada cota impuesta por el usuario), puede utilizar una técnica llamada “clipping de gradientes”.

El clipping sin desescalar resultaría en las normas/magnitudes máximas de los gradientes siendo escaladas, invalidando su umbral solicitado (que se suponía era el umbral para los gradientes sin escalar). Los gradientes contenidos en los parámetros dados del optimizador se desescalan mediante scaler.unscale (optimizer).
Puede desescalar los gradientes de los otros parámetros que se habían dado previamente a otro optimizador (como optimizer1) utilizando scaler.unscale (optimizer1). Podemos ilustrar este concepto agregando dos líneas de código:

# Desescala los gradientes de los parámetros asignados del optimizador en lugar
        scaler.unscale_(optimizer)
# Debido a que los gradientes de los parámetros asignados del optimizador están desescalados, se ajusta de forma habitual: 
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

Trabajando con Gradientes Escalados

Acumulación de gradientes

La acumulación de gradientes se basa en un concepto extremadamente básico. En lugar de actualizar los parámetros del modelo, espera y acumula los gradientes a través de lotes sucesivos para calcular la pérdida y los gradientes.

Después de un cierto número de lotes, los parámetros se actualizan dependiendo de la gradiente acumulado. Aquí hay un fragmento de código sobre cómo usar la acumulación de gradientes usando pytorch:

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            # normalizar la pérdida 
            loss = loss / iters_to_accumulate

        # Acumula gradientes escalados.
        scaler.scale(loss).backward()
          # actualización de pesos
        if (i + 1) % iters_to_accumulate == 0:
            # puede desescalar aquí si desea 
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

• La acumulación de gradientes agrega gradientes a través de un tamaño de lote adecuado de batch_per_iter * iters_to_accumulate.
  La escala debe ser calibrada para el lote efectivo; esto implica revisar las grados inf/NaN, saltarse el paso si se detectan inf/NaN, y actualizar la escala a la granularidad del lote efectivo.
  También es vital mantener los gradientes en una escala escalada y consistente en factor de escala cuando se agregan los gradientes para un lote efectivo particular.

Si las gradientes se desescalan (o el factor de escala cambia) antes de que la acumulación esté completa, la próxima pasada hacia atrás añadirá gradientes escaladas a gradientes no escaladas (o gradientes escaladas por un factor diferente), después de lo cual será imposible recuperar las gradientes no escaladas acumuladas que deben aplicarse.

• Puedes desescalar las gradientes utilizando unscale poco antes del paso, después de que se hayan acumulado todas las gradientes escaladas para el próximo paso.
  Para asegurar un lote efectivo completo, simplemente llama update al final de cada iteración donde antes llamaste step
• enumerate(data) funciona permite mantener el registro del índice del lote mientras iteramos por los datos.
• Divide la pérdida en curso por iters_to_accumulate(loss / iters_to_accumulate). Esto reduce la contribución de cada mini-batch que estamos procesando al normalizar la pérdida. Si promedias la pérdida dentro de cada lote, la división ya es correcta y no se requiere ninguna normalización adicional. Este paso puede no ser necesario dependiendo de cómo calculas la pérdida.
• Cuando utilizamos scaler.scale(loss).backward(), PyTorch acumula las gradientes escaladas y las guarda hasta que llamamos a optimizer.zero_grad().

Penalización de gradiente

Al implementar una penalización de gradiente, se utiliza torch.autograd.grad() para construir las gradientes, que se combinan para formar el valor de la penalización, y luego se agrega a la pérdida. La penalización L2 sin escalado o autocasting se muestra en el ejemplo de abajo.

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # Crea gradientes
        grad_prams = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # Calcula la parte de penalización y la agrega a la pérdida
        grad_norm = 0
        for grad in grad_prams:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # Puedes clivar gradientes aquí

        optimizer.step()

Los tensores proporcionados a torch.autograd.grad() deben ser escalados para implementar una penalización de gradiente. Es necesario desescalar los gradientes antes de combinarlos para obtener el valor de la penalización. Dado que la computación de la parte de penalización es parte de la pasada forward, debe ocurrir dentro de un contexto de autocast.
Para la misma penalización L2, así es cómo se ve:

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Realizar escalado de pérdidas para la pasada en reversa de autograd.grad, resultando #scaled_grad_prams
        scaled_grad_prams = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # Crear grad_prams antes de calcular la penalidad (grad_prams debe ser #desescalado). 
        # Debido a que ningún optimizador posee scaled_grad_prams, se utiliza división convencional en lugar de scaler.unscale_:
        inv_scaled = 1./scaler.get_scale()
        grad_prams = [p * inv_scaled for p in scaled_grad_prams]

        # Se calcula y agrega la penalidad al rendimiento. 
        with autocast():
            grad_norm = 0
            for grad in grad_prams:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        # Aplica escalado a la llamada en reversa.
        # Acumula gradientes hoja escalados apropiadamente.
        scaler.scale(loss).backward()

        # Puedes desescalar aquí 

        # step() y update() proceden como habitual.
        scaler.step(optimizer)
        scaler.update()

Trabajando con Múltiples Modelos, Pérdidas y Optimizadores

Scaler.scale debe ser llamado en cada pérdida en tu red si tienes muchas de ellas.

Si tienes muchos optimizadores en tu red, puedes ejecutar scaler.unscale en cualquiera de ellos, y debes llamar a scaler.step en cada uno de ellos. Sin embargo, scaler.update debe usarse solo una vez, después de realizar el step en todos los optimizadores usados en esta iteración:

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        with autocast():
            output1 = model1(input)
            output2 = model2(input)
            loss1 = loss_fn(2 * output1 + 3 * output2, target)
            loss2 = loss_fn(3 * output1 - 5 * output2, target)

       #Aunque retain graph no está relacionado con amp, está presente en este #ejemplo, ya que ambas llamadas backward() comparten ciertas regiones del gráfico. 
        scaler.scale(loss1).backward(retain_graph=True)
        scaler.scale(loss2).backward()

        # Si deseas ver o ajustar los gradientes de los parámetros que poseen, puedes especificar qué optimizadores reciben un desescalado explícito.
        scaler.unscale_(optimizer1)

        scaler.step(optimizer1)
        scaler.step(optimizer2)

        scaler.update()

Cada optimizador examina sus gradientes en busca de infs/NaNs y toma una decisión individual sobre si omitir o no el paso. Algunos optimizadores pueden omitir el paso, mientras que otros pueden no hacerlo. Omitir el paso ocurre solo una vez cada varias cientos de iteraciones; por lo tanto, no debería afectar la convergencia. Para modelos con múltiples optimizadores, puedes reportar el problema si observas una mala convergencia después de agregar el escalado de gradientes.

Trabajar con múltiples GPUs

Uno de los mayores problemas con los modelos de Aprendizaje Profundo es que se están haciendo demasiado grandes para entrenar en un solo GPU. Puede tardar mucho en entrenar un modelo con un solo GPU, y se necesita el entrenamiento multi-GPU para preparar los modelos lo antes posible. Un researcher bien conocido pudo reducir el periodo de entrenamiento de ImageNet de dos semanas a 18 minutos o entrenar el Transformer-XL más extenso y avanzado en dos semanas en lugar de cuatro años.

DataParallel y DistributedDataParallel

Sin comprometer la calidad, PyTorch ofrece la mejor combinación de facilidad de uso y control. nn.DataParallel y nn.parallel.DistributedDataParallel son dos características de PyTorch para distribuir el entrenamiento entre varios GPUs. Puede utilizar estos wrappers fáciles de usar y cambios para entrenar la red en varios GPUs.

DataParallel en un solo proceso

En una sola máquina, DataParallel ayuda a extendir el entrenamiento sobre varios GPUs.
Vamos a echar un vistazo más de cerca cómo DataParallel realmente funciona en la práctica.
Cuando se utiliza DataParallel para entrenar una red neuronal, se llevan a cabo las siguientes etapas:

Fuente:

• El mini-batch se divide en GPU:0.
• Dividir y distribuir mini-batch a todos los GPUs disponibles.
• Copiar el modelo a los GPUs.
• Tener lugar forward pass en todos los GPUs.
• Calcular la pérdida en relación a las salidas de la red en GPU:0, así como devolver las pérdidas a los diferentes GPUs. Los gradientes deberían calcularse en cada GPU.
• Suma de gradientes en GPU:0 y aplicar el optimizador para actualizar el modelo.

Es importante notar que las preocupaciones discutidas aquí solo se aplican a autocast. El comportamiento de GradScaler no cambia. No importa si torch.nn.DataParallel crea hilos para cada dispositivo para realizar la pasada hacia adelante. El estado de autocast se comunica en cada uno y lo siguiente funcionará:

model = Model_m()
p_model = nn.DataParallel(model)

# Establece autocast en el hilo principal
with autocast():
    # Habrá autocasting en p_model. 
    output = p_model(input)
    # loss_fn también autocast
    loss = loss_fn(output)

DistributedDataParallel, un GPU por proceso

La documentación para torch.nn.parallel.DistributedDataParallel recomienda utilizar una GPU por proceso para obtener el mejor rendimiento. En esta situación, DistributedDataParallel no inicia hilos internamente; por lo tanto, el uso de autocast y GradScaler no se ve afectado.

DistributedDataParallel, múltiples GPUs por proceso

Aquí torch.nn.parallel.DistributedDataParallel puede generar un hilo secundario para ejecutar la pasada hacia adelante en cada dispositivo. La solución es la misma: aplicar autocast como parte del método forward de tu modelo para asegurarte de que esté habilitado en los hilos secundarios.

Conclusión

En este artículo, hemos :

• Introducido Apex.
• Visto cómo funcionan los Amps.
• Visto cómo se realiza la escalada de gradientes, el recorte de gradientes, la acumulación de gradientes y la pena de gradientes.
• Visto cómo podemos trabajar con múltiples modelos, pérdidas y optimizadores.
• Visto cómo podemos realizar DataParallel en un solo proceso al trabajar con múltiples GPU.

Referencias

https://developer.nvidia.com/blog/apex-pytorch-easy-mixed-precision-training/
https://nvidia.github.io/apex/amp.html
https://discuss.pytorch.org/t/accumulating-gradients/30020
https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2