計算最佳化器的大小

你正在比較不同的最佳化器來訓練模型，並需要量化最佳化器的記憶體用量，以便客觀比較。作為測試，你已載入 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")

編輯並執行程式碼