Entrenamiento con precisión mixta en PyTorch básico

Vas a usar tipos de datos de coma flotante de baja precisión para acelerar el entrenamiento de tu modelo de traducción automática. Por ejemplo, los tipos de coma flotante de 16 bits (float16) ocupan solo la mitad que sus equivalentes de 32 bits (float32). Esto acelera los cálculos de multiplicaciones de matrices y convoluciones. Recuerda que esto implica escalar los gradientes y convertir las operaciones a coma flotante de 16 bits.

Se han precargado algunos objetos: dataset, model, dataloader y optimizer.

Este ejercicio forma parte del curso

Entrenamiento eficiente de modelos de IA con PyTorch

Ver curso

Instrucciones del ejercicio

Antes del bucle, define un escalador para los gradientes usando torch.amp.GradScaler.
En el bucle, convierte las operaciones al tipo de dato de coma flotante de 16 bits usando torch.autocast como gestor de contexto.
En el bucle, escala la pérdida y llama a .backward() para crear gradientes escalados.

ejercicio interactivo práctico

Prueba este ejercicio completando este código de ejemplo.

# Define a scaler for the gradients
scaler = torch.amp.____()
for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    # Casts operations to mixed precision
    with torch.____(device_type="cpu", dtype=torch.____):
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
    # Compute scaled gradients
    scaler.____(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

Editar y ejecutar código