Obliczanie rozmiaru optymalizatora

Analizujesz różne optymalizatory do trenowania modelu i chcesz obiektywnie porównać ich zużycie pamięci. W tym celu wczytano model DistilBERT oraz optymalizator AdamW, aby zmierzyć użycie pamięci. Napisz funkcję compute_optimizer_size, która oblicza rozmiar optymalizatora.

Optymalizator AdamW (optimizer) został zdefiniowany bezpośrednio (bez Trainer), a trenowanie zostało zakończone.

To ćwiczenie jest częścią kursu

Efektywne trenowanie modeli AI z PyTorch

Zobacz kurs

Instrukcje do ćwiczenia

Oblicz liczbę elementów i rozmiar każdego tensora (tensor) w pętli for.
Oblicz całkowity rozmiar optimizer w megabajtach.
Uzyskaj dostęp do słownika stanu optymalizatora, używając odpowiedniej metody na optimizer.state.

Interaktywne ćwiczenie praktyczne

Spróbuj tego ćwiczenia, uzupełniając ten przykładowy kod.

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

Edytuj i uruchom kod