UC CV type model optimization summary

背景

CV モデルはビジネスで一般的に使用されるモデルですが、UC クラスターにおける CV モデルには GPU 使用率の大幅な改善余地があることが確認されています。この最適化が行われない場合、レイテンシ要件を満たすために大量の GPU リソースが必要となります。リソース使用率の向上とコンピューティングリソースの節約を実現するため、これらのモデル（主に PyTorch）の最適化を検討し、より多くのリソースを消費する Detectron2 モデルでテストを実施しました。

Detectron2 はセントラルネットワークをバックボーンとして特徴抽出を行い、抽出された特徴は候補ボックス生成器に入力されて候補ボックスが生成され、候補ボックスの調整、重み除去および事前フィルタリングが実行されます。最後に roi_heads でスコアリングとソートが行われ、検出スコアとバウンディングボックスが出力されます。

モデルの計算時間は主にバックボーン部分に集中しています。Python で直接オンラインデプロイに使用されるため、最適化前のシングル GPU サービス能力は低く、シングル GPU で 12 QPS を超えるとキューイングが発生し始め、20 QPS を超えるとレイテンシが劇的に増大し、サービスが異常状態になることすらありました。サービスの安定稼働を維持するため、オンライン環境ではより多くのコンピューティングリソースを使用して QPS を向上させる必要がありました。リクエストのピーク時、サービスの QPS は 1850 前後ですが、リソースを増やして QPS を向上させるこの方法では大量のリソース浪費が発生します。ピーク時の GPU 使用率は約 20% にとどまることが観測されています。

最適化手法の検討

最適化は主に 2 つの観点から検討します。1) 単一 CUDA Stream 条件下でのモデル推論速度の向上、2) 複数 CUDA Stream の活用による GPU 使用率のさらなる向上です。

モデル推論のバックエンド最適化

バックエンド最適化の主な注目点

エンジニアリングチームとして、最適化ツールの以下の特徴に注目しています。

・最適化成功率：PyTorch は柔軟で拡張性が高く、演算子の種類も豊富なため、モデルおよび推論コードの多様性が高くなります。一方で、バックエンド演算子のほとんどは PyTorch のサブセットのみをサポートしています。このような状況下で最適化成功率をいかに最大化するかが、最も重要な考慮点です。

・最適化性能：一方では、デプロイハードウェアの能力を最大限に発揮し、ユーザーの応答時間を短縮してコストを削減できます。他方では、応答時間を維持したままより複雑なディープラーニングモデルを使用して、ビジネス精度指標を向上させることができます。

・最適化とデプロイの難易度：最適化ツールはシンプルで安定した使いやすい API を提供し、ユーザーがモデル最適化のプロセスを迅速に完了できるようにする必要があります。同時に、最適化済みモデルを既存のエンジニアリングシステムに容易にデプロイできる必要があります。

主な主流最適化手法

ネイティブ TensorRT

最適化性能の観点から、CV ドメインモデルの現在のベストプラクティスは、グラフ全体を TensorRT に変換してその高速化特性を最大限に活用することです。チームは TensorRT 高速化の豊富な経験と複数の TensorRT プラグインを蓄積しています。PyTorch モデルから TensorRT への変換には主に 2 つの方法があります。

1. PyTorch モデルを ONNX にエクスポートし、onnx2trt で TensorRT ネットワークにパースする。

2. TensorRT ネットワーク API を直接呼び出して等価なネットワークを構築し、パラメーターをロードする。

上記の方法の問題点は、TensorRT がサポートする演算子の種類が ONNX より少なく、PyTorch と ONNX 間で多くの演算子が変換できないことです。そのため、これら 2 つの方法の最適化成功率は低くなります。さらに、TensorRT プラグインの一部の実装には堅牢性の問題があり、異なるモデル（roiAlign など）で数値の差異が生じる可能性があります。

TensorRT ベースのバックエンド

PyTorch 演算子の多様性を考慮し、サブグラフアーキテクチャに基づく TensorRT ベースのバックエンドをさらに検討します。サブグラフとは、最適化前に計算グラフ内の最適化可能な演算子を選択し、これらをサブグラフに集約する手法です。各サブグラフは対応するバックエンドに変換されて最適化されます。対応するバックエンドの演算子に変換できない場合は、PyTorch に戻して実行されます。

このアーキテクチャは、デプロイ成功率を確保した前提で最大限の最適化効果を実現できます。現在、サブグラフアーキテクチャで TensorRT に接続するコミュニティの主流最適化ツールは、Torch TensorRT と PAI Blade の 2 つです。

グラフ変換アーキテクチャの面では、Torch TensorRT は TorchScript と Torch FX Module の 2 種類の入力をサポートしています。PyTorch 演算子用の多数のコンバーターを備え、コンバーター内で TensorRT ネットワーク API を呼び出して演算子を変換しネットワークを構築します。PAI Blade は TorchScript を入力としてサポートしています。サブグラフは ONNX 経由で TensorRT ネットワークに変換されます。ONNX 経由のルートはコミュニティの成熟したツールを再利用できますが、中間表現のレイヤーが追加されるためカスタマイズリンクが長くなります。直接コンバーターを記述する方法はカスタマイズリンクを短縮できますが、各演算子およびその異なる条件ごとにゼロから開発する必要があるため、開発量が大幅に増加します。現段階では、ONNX 経由のルートの方が成熟度が高いと言えます。

TensorRT 機能のカバー率では、両者ともダイナミックレンジ、低精度最適化、カスタムプラグインなどの機能を提供しています。ただし、Torch TensorRT でダイナミックレンジや INT8 量子化などの高次機能を有効化するには、入力モデルをネットワーク全体で TensorRT に変換できる必要があります。一方、PAI Blade はサブグラフアーキテクチャの下で各サブグラフに対して個別にこれらの高次機能を有効化できます。

最適化成功率と最適化性能の面では、PAI Blade は TorchScript の特徴に基づいて多数の JIT パスを開発し、最適化グラフをよりクリーンで効率的にしています。また、サブグラフの範囲が広いため、より大きな最適化効果を実現しています。ただし、ONNX という中間媒介の存在により、一部の特殊な演算子は TensorRT バックエンドに変換できません。一方で、Torch TensorRT は現段階では成熟度が低く、最適化成功率も低くなっています。

TensorRT の他にも、OneFlow や PaddlePaddle など、業界で人気の他のバックエンドも調査しました。これらのバックエンドには、ネイティブ TensorRT と同様に演算子のサポートが不完全で opset が低いといった問題があります。

各方面からの比較検討を経て、最終的にユーザビリティ、堅牢性、性能に優れた PAI Blade 最適化スキームが採用され、元のモデル実装が書き換えられてグラフ変換に対応させました。

マルチ CUDA Stream 最適化

PyTorch モデルはすべてのスレッドの CUDA 演算をデフォルトストリームに配置するため、マルチスレッドが CUDA カーネル実行で直列化され、他のストリームの実行もブロックされます。マルチスレッド内で実行される推論は互いに独立しており、実行に相互依存関係がないため、ストリームを分割することができます。一般的な単一 CUDA Stream のタイムラインは以下の通りです。

最適化プロセス

モデル推論コードの書き換え

PyTorch モデルは nn.Module レベルに「明示的な」計算グラフを含んでいません。TorchScript に変換して推論用計算グラフを取得し、推論最適化を行う必要があります。TorchScript は Python のサブセットに過ぎず、サポートされる構文が限られています。一方で、PyTorch の動的グラフ特性により、モデルの推論コードの多様性が非常に高くなります。この矛盾により、推論コードを変更せずに TorchScript をエクスポートできない場合があります。そのため、Detectron2 の推論コードを修正する必要があり、主に以下の部分を含みます。

構文とロジックの書き換え

1. 変数とパラメーターへの明示的な型アノテーションの追加

TorchScript は明示的に型アノテーションが付与されていないパラメーターをテンソルとして扱います。実際の変数がテンソルでない場合、推論プロセスで型の不一致例外が発生します。

#変更前
def forward(self, x, dim):
return torch.sum(x, dim=dim)
#変更後
def forward(self, x: torch.Tensor, dim: int):
return torch.sum(x, dim=dim)

2. Dict や List などのコンテナに格納する内容の型を一貫させる

TorchScript はコンテナオブジェクトに格納するデータ型が完全に一致していることを要求します。たとえば、List 内のすべてのデータの型は一致している必要があります。

#変更前

x = [1, "str"]
num_ data = x[0]
str_ data = x[1]
#変更後
int_ x = [1,]
str_ x = ["str",]
num_ data = int_ x[0]
str_ data = str_ x[0]

3. TorchScript がサポートしない __getattr__ などのマジックメソッドの書き換え

def init():
self.local_ modules :nn.ModuleList = []
self.local_ modules_ tmp :Dict[str, nn.Module] = {}
#self.add_ module("acb", lateral_conv)
self.local_ modules_ tmp["acb"] = lateral_ conv.cuda()
self.local_ modules = nn.ModuleDict(self.local_modules_tmp)
def forward(self, input_node):
for kn, vv in self.local_ modules.items():
if name == kn:
input_ node = vv(input_node)
#input_ node = self.__ getattr__ (name)(input_node)

4. * 展開関数の置き換え

Boxlists=list (zip (* sampled_boxes)) # for ループに置き換え

5. list、dict、NoneType の修正。list は torch.nn.ModuleList に、dict は torch.nn.ModuleDict などに置き換えます。

6. バイパス機構を活用して推論と独立したコード分岐を無視する

モデルの推論コードには、トレーニング関連と推論関連の部分が同時に存在し、さらに前処理や後処理など、モデル推論自体には属さないコードも含まれる場合があります。この部分のコードは、関連する機構を通じて直接バイパスできます。

#変更前
if self.traininig:
do_ some_ training_ thing
else:
do_ some_ eval_ thing
#変更後
if self.traininig:
assert not torch.jit.is_ scripting(), "Bypass the codes blow"
do_ some_ training_ thing # assert 以降のコードは変換されない
else:
do_ some_ eval_ thing

7. 現在、デプロイ環境ではシングル GPU 推論を使用しているため、torch.nn.SyncBatchNorm の代わりに torch.nn.BatchNorm2d を使用します。

モデルパラメーターのキー修正

モデルコードの変更後、torch.nn.Module の組織形式が変化し、state_dict のキーが変更されます。キーの不一致により、事前学習済みモデルを直接ロードできなくなります。そのため、事前学習済みモデルのキーを修正して正しくロードできるようにする必要があります。

self.linear = xxx # key : "linear"
self.linear = torch.nn.ModuleList([xxx]) # key: "linear.0","linear.1",..

グラフ分割

実際のテストでは、最適化後の GPU のエンドツーエンド性能が必ずしも最良とは限りません。さらに、GPU のパフォーマンスボトルネックが CPU に起因する場合もあります。そのため、異なるネットワークコンポーネント（バックボーン、ネック、検出ヘッド）および異なるコードロジックに対してグラフを分割し、各部分をより適切に最適化します。

変換とサークルグラフ分析

PAI Blade 最適化

import torch_ blade
import torch_ blade.tensorrt
model_ path = "./opt_model_backbone.pt"

#最適化前 TorchScript のロード
script_ model = torch.jit.load(model_path)
script_ model.cuda().eval()
...

#最適化の設定

cfg = torch_ blade.Config()
cfg.optimization_ Pipeline="TensorRT" # TensorRT バックエンドを使用
#最適化の実行
with cfg:
opt_ model = torch_ blade.optimize(script_model, False, model_inputs=img)
print("Optimize finish。")
#最適化済みモデルも TorchScript であり、関連 API を呼び出してローカルにシリアル化できる
torch.jit.save(opt_model, opt_file_name)
print("save finish。")
#モデルのロードと推論
opt_ model = torch.jit.load(opt_file_name)
#推論の実行
opt_ model(img)
サークルグラフ分析

デフォルトでは、PAI Blade のサークルグラフは少なくとも 3 つの基本 OP で構成されます。TRT グループに組み込まれない時間のかかる OP については、手動でホワイトリストを設定できます。その後、サークルグラフはその OP に特別な処理を施して最適化します。さらに、最適化後の性能が期待通りでない場合は、サークルグラフ情報を分析して、TRT グループに組み込まれていない重要なサブグラフがないかを確認できます。

マルチ CUDA Stream

最適化結果

Detectron2 モデルに 2 つの最適化を適用した結果、システムの QPS と GPU 使用率が大幅に向上しました。最適化後、約 65% のコンピューティングリソースを節約できます。

まとめ

UC のモデル推論最適化チームとして、主な目標はモデルのビジネス効果に影響を与えずに、モデル推論性能を向上させ、コンピューティングリソースの使用量を削減することです。今回の最適化では、最適化成功率、最適化性能、最適化とデプロイの難易度の 3 つの観点から最適化ツールを選定し、最終的に PAI Blade をシステムに統合しました。最終的な性能最適化効果とリソースコスト削減の観点から、当初の期待に概ね沿う結果となりました。ただし、性能最適化に終わりはありません。モデル構造の継続的な進化とモデル規模の漸次的な拡大に伴い、新たな課題が引き続き出現します。それでもなお、最適化ツールの成功率、絶対性能、使いやすさは、技術選定のコア基準であり続けています。