Entraînement en précision mixte avec PyTorch basique

Vous allez utiliser des types de nombres flottants en faible précision pour accélérer l’entraînement de votre modèle de traduction automatique. Par exemple, les types flottants sur 16 bits (float16) ne font que la moitié de la taille de leurs équivalents sur 32 bits (float32). Cela accélère les calculs des multiplications de matrices et des convolutions. Rappelez-vous que cela implique de mettre à l’échelle les gradients et de convertir les opérations en flottants 16 bits.

Certains objets ont été préchargés : dataset, model, dataloader et optimizer.

Cet exercice fait partie du cours

<cours>Entraîner efficacement des modèles d’IA avec PyTorch</cours>

Voir le cours

Instructions de l’exercice

Avant la boucle, définissez un scaler pour les gradients avec torch.amp.GradScaler.
Dans la boucle, convertissez les opérations en type flottant 16 bits à l’aide de torch.autocast comme gestionnaire de contexte.
Dans la boucle, mettez à l’échelle la perte et appelez .backward() pour créer des gradients mis à l’échelle.

Exercice interactif pratique

Essayez cet exercice en complétant ce code d’exemple.

# 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()

Modifier et exécuter le code