НачатьНачать бесплатно

Вычисление размера оптимизатора

Вы исследуете различные оптимизаторы для обучения модели и хотите объективно сравнить их потребление памяти. В качестве теста вы загрузили модель DistilBERT и оптимизатор AdamW, чтобы оценить использование памяти. Напишите функцию compute_optimizer_size для вычисления размера оптимизатора.

Оптимизатор AdamW optimizer определён напрямую (без Trainer), и обучение завершено.

Это упражнение является частью курса

Эффективное обучение моделей ИИ с PyTorch

Посмотреть курс

Инструкции к упражнению

  • Вычислите количество элементов и размер каждого тензора tensor в цикле for.
  • Вычислите общий размер optimizer в мегабайтах.
  • Получите доступ к словарю состояния оптимизатора с помощью соответствующего метода объекта optimizer.state.

Интерактивное практическое упражнение

Попробуйте выполнить это упражнение, дополнив этот пример кода.

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")
Редактировать и запускать код