Platform For AI:Blade を使用した RetinaNet (Detectron2) の最適化

制限事項

操作手順

ステップ 1：モデルのエクスポート

Detectron2 は、Facebook AI Research (FAIR) が提供する、オブジェクト検出と画像セグメンテーションのための、柔軟で拡張性、設定可能性に優れたオープンソースのトレーニングフレームワークです。このフレームワークは柔軟性が高いため、従来の方法でモデルをエクスポートすると、失敗したり、不正な結果が生成されたりする可能性があります。TorchScript のデプロイメントをサポートするために、Detectron2 は TracingAdapter と scripting_with_instances という 2 つのエクスポートメソッドを提供しています。詳細については、「Detectron2 の使用方法」をご参照ください。

import torch
import numpy as np

from torch import Tensor
from torch.testing import assert_allclose

from detectron2 import model_zoo
from detectron2.export import scripting_with_instances
from detectron2.structures import Boxes
from detectron2.data.detection_utils import read_image

# scripting_with_instances を使用して RetinaNet モデルをエクスポートします。
def load_retinanet(config_path):
    model = model_zoo.get(config_path, trained=True).eval()
    fields = {
        "pred_boxes": Boxes,
        "scores": Tensor,
        "pred_classes": Tensor,
    }
    script_model = scripting_with_instances(model, fields)
    return model, script_model

# サンプル画像をダウンロードします。
# wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg
img = read_image('./input.jpg')
img = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))

# モデルのエクスポート前後の結果を実行して比較します。
pytorch_model, script_model = load_retinanet("COCO-Detection/retinanet_R_50_FPN_3x.yaml")
with torch.no_grad():
    batched_inputs = [{"image": img.float()}]
    pred1 = pytorch_model(batched_inputs)
    pred2 = script_model(batched_inputs)

assert_allclose(pred1[0]['instances'].scores, pred2[0].scores)

ステップ 2：Blade を呼び出してモデルを最適化

1. Blade 最適化インターフェイスを呼び出します。
   blade.optimize インターフェイスを呼び出してモデルを最適化します。次のコードに例を示します。blade.optimize インターフェイスの詳細については、「PyTorch モデルの最適化」をご参照ください。

   import blade
   
   test_data = [(batched_inputs,)] # PyTorch の入力データはタプルのリストです。
   optimized_model, opt_spec, report = blade.optimize(
       script_model,  # 前のステップでエクスポートされた TorchScript モデル。
       'o1',  # Blade O1 レベルの最適化を有効にします。
       device_type='gpu', # ターゲットデバイスは GPU です。
       test_data=test_data, # 最適化とテストを支援するための特定のテストデータセット。
   )

2. 最適化レポートを出力し、モデルを保存します。
   Blade によって最適化されたモデルも TorchScript モデルです。最適化が完了したら、次のコードを使用して最適化レポートを出力し、最適化されたモデルを保存します。

   # 最適化レポートを出力します。
   print("Report: {}".format(report))
   # 最適化されたモデルを保存します。
   torch.jit.save(optimized_model, 'optimized.pt')

   以下は、出力された最適化レポートです。最適化レポートのフィールドの詳細については、「最適化レポート」をご参照ください。

   Report: {
     "software_context": [
       {
         "software": "pytorch",
         "version": "1.8.1+cu102"
       },
       {
         "software": "cuda",
         "version": "10.2.0"
       }
     ],
     "hardware_context": {
       "device_type": "gpu",
       "microarchitecture": "T4"
     },
     "user_config": "",
     "diagnosis": {
       "model": "unnamed.pt",
       "test_data_source": "user provided",
       "shape_variation": "undefined",
       "message": "Unable to deduce model inputs information (data type, shape, value range, etc.)",
       "test_data_info": "0 shape: (3, 480, 640) data type: float32"
     },
     "optimizations": [
       {
         "name": "PtTrtPassFp16",
         "status": "effective",
         "speedup": "3.77",
         "pre_run": "40.64 ms",
         "post_run": "10.78 ms"
       }
     ],
     "overall": {
       "baseline": "40.73 ms",
       "optimized": "10.76 ms",
       "speedup": "3.79"
     },
     "model_info": {
       "input_format": "torch_script"
     },
     "compatibility_list": [
       {
         "device_type": "gpu",
         "microarchitecture": "T4"
       }
     ],
     "model_sdk": {}
   }

3. 最適化前後のモデルで性能テストを実行します。
   次のコードは、性能テストの例です。

   import time
   
   @torch.no_grad()
   def benchmark(model, inp):
       for i in range(100):
           model(inp)
       torch.cuda.synchronize()
       start = time.time()
       for i in range(200):
           model(inp)
       torch.cuda.synchronize()
       elapsed_ms = (time.time() - start) * 1000
       print("Latency: {:.2f}".format(elapsed_ms / 200))
   
   # 元のモデルの速度をテストします。
   benchmark(pytorch_model, batched_inputs)
   # 最適化されたモデルの速度をテストします。
   benchmark(optimized_model, batched_inputs)

   以下は、このテストの参照結果です。

   Latency: 42.38
   Latency: 10.77

   この結果から、200 回の実行後、元のモデルと最適化されたモデルの平均レイテンシーは、それぞれ 42.38 ms と 10.77 ms であることがわかります。

ステップ 3：最適化されたモデルのロードと実行

1. オプション： 試用期間中に、認証失敗によるプログラムの予期せぬ終了を防ぐために、次の環境変数を設定します。

   export BLADE_AUTH_USE_COUNTING=1

2. PAI-Blade の使用認証を取得します。

   export BLADE_REGION=<region>
   export BLADE_TOKEN=<token>

   ビジネス要件に基づいて、次のパラメーターを設定します。

   • <region>：PAI-Blade を使用するリージョン。PAI-Blade ユーザーの DingTalk グループに参加して、PAI-Blade を使用できるリージョンを取得できます。DingTalk グループの QR コードについては、「PAI-Blade のインストール」をご参照ください。

   • <token>：PAI-Blade の使用に必要な認証トークン。PAI-Blade ユーザーの DingTalk グループに参加して、認証トークンを取得できます。DingTalk グループの QR コードについては、「PAI-Blade のインストール」をご参照ください。

3. モデルをデプロイします。
   Blade で最適化されたモデルも TorchScript モデルです。そのため、環境を切り替えることなく、最適化されたモデルをロードできます。

   import blade.runtime.torch
   import detectron2
   import torch
   
   from torch.testing import assert_allclose
   from detectron2.utils.testing import (
       get_sample_coco_image,
   )
   
   pytorch_model = model_zoo.get("COCO-Detection/retinanet_R_50_FPN_3x.yaml", trained=True).eval()
   optimized_model = torch.jit.load('optimized.pt')
   
   img = read_image('./input.jpg')
   img = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))
   
   with torch.no_grad():
       batched_inputs = [{"image": img.float()}]
       pred1 = pytorch_model(batched_inputs)
       pred2 = optimized_model(batched_inputs)
   
   assert_allclose(pred1[0]['instances'].scores, pred2[0].scores, rtol=1e-3, atol=1e-2)