计算优化器大小

您正在比较用于训练模型的不同优化器，需要量化优化器的内存占用以便客观对比。作为测试，您已加载 DistilBERT 模型并使用 AdamW 优化器，以便量化内存使用情况。请编写 compute_optimizer_size 函数来计算优化器的大小。

AdamW optimizer 已直接定义（未使用 Trainer），并且训练已完成。

本练习是课程的一部分

使用 PyTorch 高效训练 AI 模型

查看课程

练习说明

在 for 循环中计算每个 tensor 的元素个数和元素大小。
计算 optimizer 的总大小（单位：MB）。
使用 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")

编辑并运行代码