Bereken de grootte van de optimizer

Je verkent verschillende optimizers voor het trainen van een model en je moet het geheugengebruik van een optimizer kwantificeren voor een objectieve vergelijking. Als test heb je een DistilBERT-model en een AdamW-optimizer geladen zodat je het geheugengebruik kunt bepalen. Schrijf de functie compute_optimizer_size om de grootte van een optimizer te berekenen.

De AdamW-optimizer is direct gedefinieerd (zonder Trainer), en het trainen is afgerond.

Deze oefening maakt deel uit van de cursus

Efficiënt AI-modellen trainen met PyTorch

Bekijk cursus

Oefeninstructies

Bereken het aantal elementen en de grootte van elk tensor in de for-lus.
Bereken de totale grootte van de optimizer in megabytes.
Benader de optimizer state dictionary via de juiste methode op optimizer.state.

Interactieve oefening met praktijkervaring

Probeer deze oefening door deze voorbeeldcode aan te vullen.

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

Code bewerken en uitvoeren