Výpočet velikosti optimalizátoru

Zkoumáš různé optimalizátory pro trénování modelu a potřebuješ objektivně porovnat jejich paměťové nároky. Jako testovací případ sis načetl/a model DistilBERT a optimalizátor AdamW, abys mohl/a změřit využití paměti. Napiš funkci compute_optimizer_size, která vypočítá velikost optimalizátoru.

Optimalizátor AdamW optimizer byl definován přímo (bez Trainer) a trénování bylo dokončeno.

Toto cvičení je součástí kurzu

Efficient AI Model Training with PyTorch

Zobrazit kurz

Pokyny k cvičení

V cyklu for vypočítej počet prvků a velikost každého tenzoru tensor.
Vypočítej celkovou velikost optimizer v megabajtech.
Pro přístup ke stavovému slovníku optimalizátoru použij příslušnou metodu objektu optimizer.state.

Interaktivní cvičení na vyzkoušení si v praxi

Vyzkoušejte si toto cvičení dokončením tohoto ukázkového kódu.

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

Upravit a spustit kód