Calcular o tamanho do otimizador

Você está explorando diferentes otimizadores para treinar um modelo e precisa quantificar o uso de memória de um otimizador para fazer uma comparação objetiva. Como teste, você carregou um modelo DistilBERT e o otimizador AdamW para quantificar o uso da memória. Escreva a função compute_optimizer_size para calcular o tamanho de um otimizador.

O AdamW optimizer foi definido e o treinamento foi concluído com o uso do optimizer.

Este exercício faz parte do curso

Treinamento eficiente de modelos de IA com PyTorch

Ver curso

Instruções do exercício

Calcule o número de elementos e o tamanho de cada tensor no loop for.
Calcule o tamanho total do site optimizer em megabytes.
Passe no estado optimizer para compute_optimizer_size.

Exercício interativo prático

Experimente este exercício completando este código de exemplo.

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 e executar o código