1. Learn
    2. /
  2. Courses
    3. /
  3. Efficient AI Model Training with PyTorch
Connected

Exercise

Compute the optimizer size

You're exploring different optimizers for training a model, and you need to quantify an optimizer's memory usage for an objective comparison. As a test, you've loaded a DistilBERT model and AdamW optimizer so that you quantify memory usage. Write the compute_optimizer_size function to compute the size of an optimizer.

The AdamW optimizer has been defined, and training has completed using optimizer.

Instructions

100 XP
  • Compute number of elements and size of each tensor in the for loop.
  • Compute the total size of the optimizer in megabytes.
  • Pass in the optimizer state to compute_optimizer_size.