Berechne die Optimierergröße

Du untersuchst verschiedene Optimierer für das Training eines Modells und musst die Speichernutzung eines Optimierers quantifizieren, um objektiv vergleichen zu können. Als Test hast du ein DistilBERT-Modell und den AdamW-Optimierer geladen, damit du die Speichernutzung messen kannst. Schreibe die Funktion compute_optimizer_size, um die Größe eines Optimierers zu berechnen.

Der AdamW-optimizer wurde definiert, und das Training wurde mit optimizer abgeschlossen.

Diese Übung ist Teil des Kurses

Effizientes KI-Modelltraining mit PyTorch

Kurs anzeigen

Anleitung zur Übung

Berechne die Anzahl der Elemente und die Größe jedes tensor in der for-Schleife.
Berechne die Gesamtgröße des optimizer in Megabyte.
Übergib den optimizer-State an compute_optimizer_size.

Interaktive Übung

Vervollständige den Beispielcode, um diese Übung erfolgreich abzuschließen.

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(____.____.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")

Code bearbeiten und ausführen