Calculer la taille de l’optimiseur

Vous explorez différents optimiseurs pour entraîner un modèle et vous devez quantifier l’usage mémoire d’un optimiseur pour pouvoir comparer objectivement. À titre de test, vous avez chargé un modèle DistilBERT et un optimiseur AdamW afin de mesurer la mémoire utilisée. Écrivez la fonction compute_optimizer_size pour calculer la taille d’un optimiseur.

L’optimizer AdamW a été défini directement (sans Trainer), et l’entraînement est terminé.

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

Calculez le nombre d’éléments et la taille de chaque tensor dans la boucle for.
Calculez la taille totale de l’optimizer en mégaoctets.
Accédez au dictionnaire d’état de l’optimiseur en utilisant la méthode appropriée sur optimizer.state.

Exercice interactif pratique

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

def compute_optimizer_size(optimizer_state):
    total_size_megabytes, total_num_elements = 0, 0
    for params in optimizer_state:
        for name, tensor in params.items():
            tensor = torch.tensor(tensor)
            # Compute number of elements and size of each tensor
            num_elements, element_size = tensor.____(), tensor.____()
            total_num_elements += num_elements
            # Compute the total size in megabytes
            total_size_megabytes += ____ * ____ / (1024 ** 2)
    return total_size_megabytes, total_num_elements

# Pass in the optimizer state
total_size_megabytes, total_num_elements = compute_optimizer_size(optimizer.state.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")

Modifier et exécuter le code