オプティマイザのサイズを計算する

モデルの学習に使用するオプティマイザを比較するために、各オプティマイザのメモリ使用量を定量的に評価します。今回は、DistilBERT モデルと AdamW オプティマイザを読み込んだ状態で、メモリ使用量を計測します。オプティマイザのサイズを計算する compute_optimizer_size 関数を作成しましょう。

AdamW の optimizer は（Trainer を使わずに）直接定義されており、学習はすでに完了しています。

この演習はコースの一部です

PyTorch による効率的な AI モデルトレーニング

コースを見る

演習の手順

for ループ内で、各 tensor の要素数と各要素のサイズを計算します。
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")

コードを編集して実行