Calculează dimensiunea optimizatorului

Explorezi diferiți optimizatori pentru antrenarea unui model și trebuie să cuantifici consumul de memorie al unui optimizator, în vederea unei comparații obiective. Ca test, ai încărcat un model DistilBERT și un optimizator AdamW pentru a măsura utilizarea memoriei. Scrie funcția compute_optimizer_size care să calculeze dimensiunea unui optimizator.

Optimizatorul AdamW optimizer a fost definit direct (fără Trainer), iar antrenamentul s-a finalizat.

Acest exercițiu face parte din cursul

Antrenament eficient al modelelor AI cu PyTorch

Vezi cursul

Instrucțiuni pentru exercițiu

Calculează numărul de elemente și dimensiunea fiecărui tensor din bucla for.
Calculează dimensiunea totală a optimizer-ului în megaocteți.
Accesează dicționarul de stare al optimizatorului folosind metoda potrivită pe optimizer.state.

Exercițiu interactiv practic

Încearcă acest exercițiu completând acest cod de exemplu.

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

Editează și rulează codul