Calcula el tamaño del optimizador

Estás explorando distintos optimizadores para entrenar un modelo y necesitas cuantificar el uso de memoria de cada optimizador para poder compararlos objetivamente. Como prueba, has cargado un modelo DistilBERT y un optimizador AdamW para medir el uso de memoria. Escribe la función compute_optimizer_size para calcular el tamaño de un optimizador.

El optimizer AdamW se ha definido directamente (sin Trainer), y el entrenamiento ya ha finalizado.

Este ejercicio forma parte del curso

Entrenamiento eficiente de modelos de IA con PyTorch

Ver curso

Instrucciones del ejercicio

Calcula el número de elementos y el tamaño de cada tensor en el bucle for.
Calcula el tamaño total del optimizer en megabytes.
Accede al diccionario de estado del optimizador usando el método adecuado sobre optimizer.state.

ejercicio interactivo práctico

Prueba este ejercicio completando este código de ejemplo.

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

Editar y ejecutar código