Calculer la taille de l’optimiseur

Vous comparez différents optimiseurs pour entraîner un modèle et vous devez quantifier l’utilisation mémoire de chaque optimiseur pour une comparaison objective. À titre de test, vous avez chargé un modèle DistilBERT et un optimiseur AdamW afin de mesurer l’empreinte mémoire. Écrivez la fonction compute_optimizer_size pour calculer la taille d’un optimiseur.

L’optimizer AdamW a été défini et l’entraînement a été effectué avec optimizer.

Cet exercice fait partie du cours

Entraîner efficacement des modèles d’IA avec PyTorch

Afficher le cours

Instructions

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.
Passez l’état de l’optimizer à compute_optimizer_size.

Exercice interactif pratique

Essayez cet exercice en complétant cet exemple de code.

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(____.____.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")

Modifier et exécuter le code