Calcula el tamaño del optimizador

Estás explorando distintos optimizadores para entrenar un modelo, y necesitas cuantificar el uso de memoria de un optimizador para realizar una comparación objetiva. Como prueba, has cargado un modelo DistilBERT y el optimizador AdamW para que cuantifiques el uso de memoria. Escribe la función compute_optimizer_size para calcular el tamaño de un optimizador.

Se ha definido el AdamW optimizer y se ha completado el entrenamiento utilizando optimizer.

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 de optimizer en megabytes.
Pasa en el estado optimizer a compute_optimizer_size.

Ejercicio interactivo práctico

Prueba este ejercicio completando el código de muestra.

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

Editar y ejecutar código