Calculer la taille de l'optimiseur

Vous étudiez différents optimiseurs pour la formation d'un modèle et vous devez quantifier l'utilisation de la mémoire d'un optimiseur pour une comparaison objective. A titre de test, vous avez chargé un modèle DistilBERT et l'optimiseur AdamW afin de quantifier l'utilisation de la mémoire. Écrivez la fonction compute_optimizer_size pour calculer la taille d'un optimiseur.

L'AdamW optimizer a été défini et la formation s'est achevée en utilisant optimizer.

Cet exercice fait partie du cours

Entraînement efficace de modèles d'IA avec PyTorch

Afficher le cours

Instructions

Calculez le nombre d'éléments et la taille de chaque site tensor dans la boucle for.
Calculez la taille totale du site optimizer en mégaoctets.
Passez dans l'état 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