Calcola la dimensione dell'optimizer

Stai valutando diversi optimizer per addestrare un modello e hai bisogno di quantificare l'uso di memoria di ciascuno per un confronto oggettivo. Come prova, hai caricato un modello DistilBERT e l'optimizer AdamW per poter misurare l'utilizzo di memoria. Scrivi la funzione compute_optimizer_size per calcolare la dimensione di un optimizer.

L'optimizer AdamW è stato definito direttamente (senza Trainer) e l'addestramento è terminato.

Questo esercizio fa parte del corso

Efficient AI Model Training with PyTorch

Visualizza corso

Istruzioni dell'esercizio

Calcola il numero di elementi e la dimensione di ogni tensor nel ciclo for.
Calcola la dimensione totale dell'optimizer in megabyte.
Accedi al dizionario di stato dell'optimizer usando il metodo appropriato su optimizer.state.

esercizio interattivo pratico

Prova questo esercizio completando questo codice di esempio.

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")

Modifica ed esegui il codice