Platform For AI:Blade とカスタム C++ 演算子による RetinaNet の最適化

制限事項

操作手順

ステップ 1：カスタム C++ 演算子を含む PyTorch モデルの作成

Blade は PyTorch の TorchScript 拡張機構とシームレスに連携します。本セクションでは、TorchScript 拡張機能を用いて RetinaNet の後処理部分を実装する方法について説明します。TorchScript カスタム演算子の詳細については、「EXTENDING TORCHSCRIPT WITH CUSTOM C++ OPERATORS」をご参照ください。本トピックで説明する RetinaNet 後処理のプログラムロジックは、NVIDIA オープンソースコミュニティから提供されたものです。詳細については、「Retinanet-Examples」をご参照ください。本トピックでは、コアコードを抽出して、カスタム演算子の開発・実装方法を解説します。

1. サンプルコードをダウンロードして解凍します。

   wget -nv https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/tutorials/retinanet_example/retinanet-examples.tar.gz -O retinanet-examples.tar.gz
   tar xvfz retinanet-examples.tar.gz 1>/dev/null

2. カスタム C++ 演算子をコンパイルします。

   公式 PyTorch ドキュメントでは、カスタム演算子のコンパイル方法として、CMake を用いたビルド、JIT（Just-in-Time）コンパイル、Setuptools を用いたビルドの 3 つが紹介されています。「EXTENDING TORCHSCRIPT WITH CUSTOM C++ OPERATORS」をご参照ください。これらの方法はそれぞれ異なるシナリオに適していますので、ご要件に応じて選択してください。本トピックでは簡便性を重視し、JIT コンパイル方式を採用します。以下に例を示します。

   import torch.utils.cpp_extension
   import os
   codebase="retinanet-examples"
   sources=['csrc/extensions.cpp',
            'csrc/cuda/decode.cu',
            'csrc/cuda/nms.cu',]
   sources = [os.path.join(codebase,src) for src in sources]
   torch.utils.cpp_extension.load(
       name="custom",
       sources=sources,
       build_directory=codebase,
       extra_include_paths=['/usr/local/TensorRT/include/', '/usr/local/cuda/include/', '/usr/local/cuda/include/thrust/system/cuda/detail'],
       extra_cflags=['-std=c++14', '-O2', '-Wall'],
       extra_cuda_cflags=[
           '-std=c++14', '--expt-extended-lambda',
           '--use_fast_math', '-Xcompiler', '-Wall,-fno-gnu-unique',
           '-gencode=arch=compute_75,code=sm_75',],
       is_python_module=False,
       with_cuda=True,
       verbose=False,
   )

   プログラム実行後、コンパイル済みの custom.so ファイルが retinanet-examples フォルダに保存されます。

3. RetinaNet の後処理部分をカスタム C++ 演算子に置き換えます。

   簡便性を考慮し、adapter_forward を RetinaNet.forward に直接置き換えることができます。adapter_forward 関数では、decode_cuda および nms_cuda のカスタム C++ 演算子を用いて RetinaNet の後処理部分を実装しています。以下に例を示します。

   import os
   import torch
   from typing import Tuple, Dict, List, Optional
   codebase="retinanet-examples"
   torch.ops.load_library(os.path.join(codebase, 'custom.so'))
   
   decode_cuda = torch.ops.retinanet.decode
   nms_cuda = torch.ops.retinanet.nms
   
   # この関数の主なコードは RetinaNet.forward と同じですが、後処理部分は decode_cuda および nms_cuda を用いて実装されています。
   def adapter_forward(self, batched_inputs: Tuple[Dict[str, torch.Tensor]]):
       images = self.preprocess_image(batched_inputs)
       features = self.backbone(images.tensor)
       features = [features[f] for f in self.head_in_features]
       cls_heads, box_heads = self.head(features)
       cls_heads = [cls.sigmoid() for cls in cls_heads]
       box_heads = [b.contiguous() for b in box_heads]
   
       # 後処理部分。
       strides = [images.tensor.shape[-1] // cls_head.shape[-1] for cls_head in cls_heads]
       decoded = [
           decode_cuda(
               cls_head,
               box_head,
               anchor.view(-1),
               stride,
               self.test_score_thresh,
               self.test_topk_candidates,
           )
           for stride, cls_head, box_head, anchor in zip(
               strides, cls_heads, box_heads, self.cell_anchors
           )
       ]
   
       # 非極大抑制（NMS）部分。
       decoded = [torch.cat(tensors, 1) for tensors in zip(decoded[0], decoded[1], decoded[2])]
       return nms_cuda(decoded[0], decoded[1], decoded[2], self.test_nms_thresh, self.max_detections_per_image)
   
   from detectron2.modeling.meta_arch import retinanet
   
   # RetinaNet.forward を adapter_forward に置き換えます。
   retinanet.RetinaNet.forward = adapter_forward

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

Detectron2 は、Facebook AI Research（FAIR）が開発した、オブジェクト検出および画像セグメンテーションのトレーニング向けの柔軟性、拡張性、および設定可能性に優れたオープンソースのフレームワークです。その柔軟性ゆえに、従来の方法でモデルをエクスポートすると、失敗するか、不正確な結果が得られる場合があります。TorchScript デプロイメントをサポートするために、Detectron2 では、TracingAdapter および scripting_with_instances の 2 つのエクスポート方法が提供されています。詳細については、「Detectron2 Usage」をご参照ください。

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()
    # PyTorch モデルに新しい cell_anchors 属性を設定します。
    model.cell_anchors = [c.contiguous() for c in model.anchor_generator.cell_anchors]
    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], pred2[0])

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

1. Blade 最適化 API を呼び出します。

   blade.optimize API を呼び出してモデルを最適化できます。以下に例を示します。blade.optimize API の詳細については、「PyTorch モデルの最適化」をご参照ください。

   import os
   import blade
   import torch
   
   # カスタム C++ 演算子のダイナミックリンクライブラリを読み込みます。
   codebase="retinanet-examples"
   torch.ops.load_library(os.path.join(codebase, 'custom.so'))
   
   blade_config = blade.Config()
   blade_config.gpu_config.disable_fp16_accuracy_check = True
   
   test_data = [(batched_inputs,)] # PyTorch の入力データはタプルのリストです。
   
   with blade_config:
       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(script_model, 'script_model.pt')
   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.92",
         "pre_run": "40.72 ms",
         "post_run": "10.39 ms"
       }
     ],
     "overall": {
       "baseline": "40.64 ms",
       "optimized": "10.41 ms",
       "speedup": "3.90"
     },
     "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(script_model, batched_inputs)
   # 最適化後のモデルの速度をテストします。
   benchmark(optimized_model, batched_inputs)

   このテストの参考結果を以下に示します。

   Latency: 40.65
   Latency: 10.46

   結果より、200 回の実行における平均レイテンシーは、元のモデルで 40.65 ms、最適化済みモデルで 10.46 ms であることが確認できます。

ステップ 4：最適化済みモデルの読み込みと実行

1. （任意） 試用期間中は、認証失敗による予期せぬプログラム終了を防ぐため、以下の環境変数設定を追加してください。

   export BLADE_AUTH_USE_COUNTING=1

2. PAI-Blade の利用認証を行います。

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

   以下のパラメーターを業務要件に応じて設定してください：

   • <region>：PAI-Blade を利用するリージョンです。利用可能なリージョンは、PAI-Blade ユーザーの DingTalk グループに参加することで確認できます。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
   import numpy as np
   import os
   from detectron2.data.detection_utils import read_image
   from torch.testing import assert_allclose
   
   # カスタム C++ 演算子のダイナミックリンクライブラリを読み込みます。
   codebase="retinanet-examples"
   torch.ops.load_library(os.path.join(codebase, 'custom.so'))
   
   script_model = torch.jit.load('script_model.pt')
   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 = script_model(batched_inputs)
       pred2 = optimized_model(batched_inputs)
   
   assert_allclose(pred1[0], pred2[0], rtol=1e-3, atol=1e-2)