Обчисліть розмір оптимізатора
Ви досліджуєте різні оптимізатори для тренування моделі й хочете кількісно оцінити використання пам’яті оптимізатором для об’єктивного порівняння. Як тест ви завантажили модель 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")