Größe des Optimierers berechnen

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

Der AdamW-optimizer wurde direkt (ohne Trainer) definiert, und das Training ist abgeschlossen.

Diese Übung ist Teil des Kurses

<Kurs>Effizientes KI-Modelltraining mit PyTorch</Kurs>

Kurs ansehen

Übungsanweisungen

Ermittle Anzahl der Elemente und Größe jedes tensor in der for-Schleife.
Berechne die Gesamtgröße des optimizer in Megabyte.
Greife über die passende Methode auf optimizer.state auf das Zustands-Dictionary des Optimierers zu.

Interaktive praktische Übung

Versuche dich an dieser Übung, indem du diesen Beispielcode vervollständigst.

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 bearbeiten und ausführen