Platform For AI:PyTorch モデルの最適化

ステップ 1：依存関係のインポート

import os
import time
import torch
import torchvision.models as models
import blade

ステップ 2：モデルの読み込みと変換

PAI-Blade は torch.jit.ScriptModule オブジェクトのみを受け付けます。標準の nn.Module モデルはサポートされません。PAI-Blade に渡す前に、torch.jit.script を使用してモデルを変換してください。

model = models.resnet50().float().cuda()  # ResNet50 モデルを読み込みます。
model = torch.jit.script(model).eval()   # ScriptModule に変換します。
dummy = torch.rand(1, 3, 224, 224).cuda() # テスト用入力テンソルを作成します。

ステップ 3：最適化の実行

blade.optimize を呼び出してモデルを最適化します。このメソッドは以下の主要なパラメーターを受け付けます。

全パラメーターのリファレンスについては、「Python メソッド」をご参照ください。

optimized_model, opt_spec, report = blade.optimize(
    model,                 # 最適化対象のモデル。
    'o1',                  # 最適化レベル。
    device_type='gpu',     # 対象デバイスタイプ。
    test_data=[(dummy,)],  # テスト用入力データ。
)

blade.optimize は以下の 3 つのオブジェクトを返します。

• `optimized_model`：推論に即時使用可能な、最適化済みの torch.jit.ScriptModule オブジェクト。

• `opt_spec`：最適化の再現に必要な外部依存関係（構成、環境変数、リソースファイル）。これらの依存関係を有効化するには、Python の with 文を使用します。

• `report`：最適化レポート。直接出力可能です。各フィールドの詳細については、「最適化レポート」をご参照ください。

最適化には時間がかかります。 PAI-Blade は、最適化済みモデルを生成する前に複数の分析フェーズを実行します。最適化の進行状況は以下のように表示されます。

[Progress] 5%, phase: user_test_data_validation.
[Progress] 10%, phase: test_data_deduction.
[Progress] 15%, phase: CombinedSwitch_4.
[Progress] 95%, phase: model_collecting.

このコンパイルは 1 回限りのコストです。最適化が完了すると、optimized_model を使用した推論は、元のモデルよりも一貫して高速になります。

ステップ 4：最適化レポートの確認

print("Report: {}".format(report))

レポートには、適用された最適化とそのレイテンシーへの影響が示されます。サンプルレポートを以下に示します。

{
  "optimizations": [
    {
      "name": "PtTrtPassFp32",
      "status": "effective",
      "speedup": "1.50",     // この最適化による速度向上率。
      "pre_run": "5.29 ms",  // 最適化前のレイテンシー。
      "post_run": "3.54 ms"  // 最適化後のレイテンシー。
    }
  ],
  "overall": {
    "baseline": "5.30 ms",    // 元のモデルのレイテンシー。
    "optimized": "3.59 ms",   // 最適化済みモデルのレイテンシー。
    "speedup": "1.48"         // 端から端までの速度向上率。
  }
}

この例では、端から端までの速度向上率は 1.48 倍 であり、最適化済みモデルの実行時間は 3.59 ms で、元のモデルの 5.30 ms より短くなっています。「"status": "effective"」を確認することで、測定可能な効果をもたらした最適化パスを特定できます。

ステップ 5：推論パフォーマンスのベンチマーク

以下の関数を使用して、元のモデルおよび最適化済みモデルの定常状態レイテンシーを測定します。最初の 100 回の実行は GPU のウォームアップに使用し、その後の 200 回の実行が計測対象となります。

@torch.no_grad()
def benchmark(model, inp):
    for i in range(100):    # ウォームアップ。
        model(inp)
    start = time.time()
    for i in range(200):    # 計測対象の実行。
        model(inp)
    elapsed_ms = (time.time() - start) * 1000
    print("Latency: {:.2f}".format(elapsed_ms / 200))

benchmark(model, dummy)           # 元のモデル。
benchmark(optimized_model, dummy) # 最適化済みモデル。