PyTorchモデルを最適化する

AI (PAI) 向け機械学習プラットフォーム-ブレードを使用すると、さまざまな方法でモデルを最適化できます。ホイールパッケージは、ローカル環境にのみインストールする必要があります。次に、Pythonメソッドを呼び出してモデルを最適化できます。このトピックでは、PAI-Bladeを使用してPyTorchモデルを最適化する方法について説明します。この例では、NVIDIA Tesla T4 GPUが使用されます。

前提条件

PyTorchがインストールされています。 PAI-Bladeのホイールパッケージが取り付けられています。
PyTorchモデルは訓練されています。この例では、オープンResNet50モデルが使用されます。

PAI-Bladeおよびその他の依存関係ライブラリをインポートします。
```
import os
import time
import torch
import torchvision.models as models
import blade
```
torchvisionライブラリからResNet50モデルを読み込みます。 PAI-BladeはScriptModulesのみをサポートします。したがって、ResNet50モデルをScriptModuleに変換する必要があります。
```
model = models.resnet50().float().cuda()  # Prepare the model. 
model = torch.jit.script(model).eval()    # Convert the model into a ScriptModule. 
dummy = torch.rand(1, 3, 224, 224).cuda() # Construct test data. 
```
ResNet50モデルを最適化するには、blade.optimizeメソッドを呼び出します。次のサンプルコードは、モデルを最適化する方法の例を示しています。モデルの最適化中に質問がある場合は、PAI-BladeユーザーのDingTalkグループに参加し、テクニカルサポートスタッフに相談してください。
```
optimized_model, opt_spec, report = blade.optimize(
    model,                 # The model to be optimized. 
    'o1',                  # The optimization level. Valid values: o1 and o2. 
    device_type='gpu',     # The type of the device on which the model is run. Valid values: gpu and cpu. 
    test_data=[(dummy,)],  # The test data. The test data used for a PyTorch model is a list of tuples of tensors. 
)
```
blade.optimizeメソッドは、次のオブジェクトを返します。
- optimized_model: 最適化されたモデル。この例では、torch.jit.ScriptModuleオブジェクトが返されます。
- opt _ spec: 最適化の結果を再現するために必要な外部依存する. 外部依存関係には、構成情報、環境変数、およびリソースファイルが含まれます。 PythonでWITHステートメントを実行して、外部依存関係を有効にすることができます。
- report: 最適化レポート。直接表示できます。最適化レポートのパラメーターの詳細については、「最適化レポート」をご参照ください。
モデル最適化中に、最適化の進行状況が表示されます。次のサンプルコードは、例を示します。
```
[Progress] 5%, phase: user_test_data_validation.
[Progress] 10%, phase: test_data_deduction.
[Progress] 15%, phase: CombinedSwitch_4.
[Progress] 95%, phase: model_collecting.
```

最適化レポートを表示します。

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

最適化レポートでは、最適化効果を実現する最適化項目を表示できます。次のサンプルコードに例を示します。

Report: {
  // ......
  "optimizations": [
    {
      "name": "PtTrtPassFp32",
      "status": "effective",
      "speedup": "1.50",     // The acceleration ratio. 
      "pre_run": "5.29 ms",  // The latency before acceleration. 
      "post_run": "3.54 ms"  // The latency after acceleration. 
    }
  ],
  // The end-to-end optimization results. 
  "overall": {
    "baseline": "5.30 ms",    // The latency of the original model. 
    "optimized": "3.59 ms",   // The latency of the optimized model. 
    "speedup": "1.48"         // The acceleration ratio. 
  },
  // ......
}

モデル最適化の前後のパフォーマンスを比較します。

@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))

# Measure the speed of the original model. 
benchmark(model, dummy)

# Measure the speed of the optimized model. 
benchmark(optimized_model, dummy)

拡張情報

blade.optimizeメソッドを呼び出すと、modelパラメーターに最適化するモデルを複数の方法で指定できます。 PyTorchモデルを最適化するには、次のいずれかの方法でモデルを指定します。

torch.jit.ScriptModuleオブジェクトを指定します。
torch.jit.saveメソッドを使用して保存されたモデルファイルからtorch.jit.ScriptModuleオブジェクトを読み込みます。

この例では、メモリ内のtorch.jit.ScriptModuleオブジェクトがblade.optimizeメソッドに指定されています。次のサンプルコードは、モデルファイルからモデルを読み込む方法の例を示しています:

optimized_model, opt_spec, report = blade.optimize(
    'path/to/torch_model.pt',
    'o1',
    device_type='gpu'
)

次のステップ

PAI-Bladeを使用してモデルを最適化した後、Pythonで最適化モデルを実行するか、PAIのElastic Algorithm service (EAS) で最適化モデルをサービスとしてデプロイできます。 PAI-Bladeは、最適化されたモデルを独自のアプリケーションに統合するのに役立つC ++ 用のSDKも提供しています。