Beräkna optimerarens storlek

Du utforskar olika optimerare för att träna en modell och behöver kvantifiera minnesanvändningen för en objektiv jämförelse. Som test har du läst in en DistilBERT-modell och en AdamW-optimerare för att mäta minnesanvändningen. Skriv funktionen compute_optimizer_size för att beräkna storleken på en optimerare.

AdamW-optimeraren optimizer har definierats direkt (utan Trainer) och träningen är avslutad.

Den här övningen är en del av kursen

Effektiv AI-modellträning med PyTorch

Visa kurs

Övningsinstruktioner

Beräkna antalet element och storleken på varje tensor i for-loopen.
Beräkna den totala storleken på optimizer i megabyte.
Hämta optimerarens tillståndsdictionary med lämplig metod på optimizer.state.

Interaktiv övning med praktiskt arbete

Testa den här övningen genom att slutföra den här exempelkoden.

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

Redigera och kör kod