ハードコアの詳細な Segment Anything Model (SAM) TensorRT モデル変換

まとめ

Segment Anything Model (SAM) は、Facebook Research によって最近オープンソース化された新しい画像セグメンテーション タスクおよびモデルです。このモデルは、プロンプトになるように設計およびトレーニングされているため、ゼロショット転送ゼロ サンプルを新しい画像分布およびタスクに転送できます。そのセグメンテーション効果は素晴らしく、SOTA をセグメント化するための現在のアルゴリズムです。アルゴリズムの詳細についてはインターネット上に多くの説明がありますが、この記事では主にモデルを TensorRT モデルに変換して、後のデプロイメントを容易にし、モデル推論を高速化する方法を説明します。

SAMコード

SAM公式サイト

この記事の TensorRT コード

簡単な紹介

セグメント エニシング モデル (SAM) モデルは、図 1 に示すように、イメージ エンコーダー、ヒント エンコーダー、およびマスク デコーダーの3 つのコンポーネントで構成されます。

図 1: セグメンテーション エブリシング モデル (SAM) の概要。ヘビーウェイト画像エンコーダは画像埋め込みを出力します。これはさまざまな入力キューによって効率的にクエリされ、償却されたリアルタイム速度でオブジェクト マスクを生成します。複数のオブジェクトに対応するあいまいなキューの場合、SAM は複数の効果的なマスクと関連する信頼スコアを出力できます。

画像エンコーダ 画像エンコーダ。スケーラビリティと強力な事前トレーニング方法にインスピレーションを得て、高解像度の入力を処理するために最小限に適合した MAE の事前トレーニング済みビジョン トランスフォーマー (ViT) を使用します。画像エンコーダーは画像ごとに 1 回実行され、モデルにヒントを与える前に適用できます。

プロンプトエンコーダ プロンプトエンコーダ。スパース (ポイント、ボックス、テキスト) と密 (マスク) の 2 つのキュー セットを考えてみましょう。ポイントとボックスは、キュー タイプごとに学習されたエンベディングを合計した位置エンコーディングによって表され、フリー テキストは CLIP の既製のテキスト エンコーダーを使用して表されます。密なキュー (つまりマスク) は畳み込み埋め込みを使用し、画像埋め込み要素と加算されます。

マスク デコーダ マスク デコーダ。マスク デコーダは、画像の埋め込み、プロンプトの埋め込み、および出力トークンをマスクに効率的にマッピングすることでこれを行います。この設計では、変更された Transformer デコーダ ブロックと、その後に動的マスク予測ヘッダーが続きます。修正されたデコーダー ブロックは、キューのセルフ アテンションとクロス アテンション (ヒントから画像への埋め込み、またはその逆) の両方向を使用して、すべての埋め込みを更新します。両方のブロックを実行した後、画像埋め込みをアップサンプリングし、MLP が出力トークンを動的線形分類器にマップし、画像の各位置でマスクされた前景確率を計算します。

モデル変換プロセス

効率を上げてモデルの高速化を図る必要があるのですが、具体的にはどうすればよいのでしょうか？ディープラーニングモデルを高速化するために現在一般的に使われている手法は、pytorch/tensorflowなどで表されるモデルをTensorRTで表されるモデルに変換することです。TensorRT は NVIDIA によって開発されたフレームワークで、モデル推論を高速化できます。実際、テスト段階でトレーニングされたモデルの速度が向上します。たとえば、モデルが 50 ミリ秒の速度で画像をテストする場合、tensorRT を使用すると、加速するには 10ms しかかからないかもしれません。TensorRT の詳細については、この記事では詳しく説明しません。公式 Web サイトをご自身で参照してください。深層学習モデルの全体的な高速化を 2 つの部分に分けました。

1. モデル変換部分。Pytorch/Tensorflowモデル -> TensorRTモデルの変換を実現します。
2. モデル推論 (推論) セクション。TensorRT モデルを使用してモデルを推論します。

Pytorch モデルから対応する TensorRT モデルを取得するにはどうすればよいですか? 一般に次の 2 つの方法があります。

1. **「torch2trt 」** による変換;
2. 「Pytorch -> onnx -> TensorRT」。このパスは最も広く使用されており、まず Pytorch モデルを ONNX で表されるモデルに変換し、次に ONNX で表されるモデルを TensorRT で表されるモデルに変換します。この方法は、本稿で焦点を当てる方法でもあります。

Pytorch モデルから ONNX モデルへ

Pytorch -> ONNX 変換は、Pytorch の組み込み API を使用することで比較的簡単です。

torch.onnx.export(model,
                  x,
                  "./ckpts/onnx_models/{}.onnx".format(model_name),
                  input_names=input_names,
                  output_names=output_names,
                  opset_version=16)

ここで強調する必要があるのは、**「opset_version」** パラメーターです。公式の onnx はまだ更新中であるため、一部の pytorch 演算子のみが変換でき、かなりの数の演算子は変換できません。したがって、変換するときは、より多くの演算子を変換できるように、opset_version の最新バージョンを選択しようとします。現在、ONNX はオペレーターと対応するバージョンを正式にサポートしています。

ONNX モデルを TensorRT モデルに変換する

ONNX -> TensorRT を変換する前に、 onnx-simplifierツールを使用して変換された ONNX モデルを単純化することを強くお勧めします。そうしないと、次の変換でエラーが報告される可能性があります。onnx-simplifier は ONNX モデルを単純化するツールです. 先ほど変換した ONNX モデルは実際には非常に冗長です. 一部の操作 (IF 判定など) は必要ありません, それらの冗長な部分は次のとおりです ONNX->TensorRT の変換モデルは不要なエラーを引き起こし、モデルのメモリを増加させる可能性があるため、モデルを簡素化する必要があります。

次に、ONNX モデルを TensorRT モデルに変換する必要があります。まず、NVIDIA 公式 Web サイトから TensorRT-8.6.1.6 ツールキットをダウンロードし、Ubuntu システムのユーザー ルート ディレクトリで解凍し、環境変数を設定する必要があります。 **モデル変換には「公式ツール trtexec」**をご利用ください。このツールは、前にダウンロードした TensorRT フォルダーに既に存在します。

# 在python环境中安装TensorRT包
pip install ~/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp38-none-linux_x86_64.whl

# 设置环境变量
export PATH=$HOME/TensorRT-8.6.1.6/targets/x86_64-linux-gnu/bin:$PATH
export TENSORRT_DIR=$HOME/TensorRT-8.6.1.6:$TENSORRT_DIR
export LD_LIBRARY_PATH=$HOME/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH


#输入命令
./trtexec --onnx=pytorch.onnx --saveEngine=pytorch.engine --workspace=4096

エラーが報告されない場合は、変換された TensorRT モデルであるpytorch.engineという名前のモデルを取得します。この時点で、モデル変換部分はすべて終了です。

モデル変換

SAM モデルの 3 つのコンポーネントを大まかに紹介した後、次のステップはモデル変換の本題に入ります。画像埋め込みモデルは主に特徴抽出に VIT を使用し、このステップは 1 回だけ実行されるため、このモジュールのモデルは次のようになります。プロンプト エンコーダー モデルとマスク デコーダー モデルはモデル変換のために結合されます。

画像埋め込みモデルをonnxに変換

python scripts/onnx2trt.py --img_pt2onnx --sam_checkpoint weights/sam_vit_h_4b8939.pth --model_type default

画像埋め込みモジュール onnx モデルを TensorRT モデルに変換

trtexec --onnx=embedding_onnx/sam_default_embedding.onnx --workspace=4096 --saveEngine=weights/sam_default_embedding.engine

ここまでで画像埋め込みモジュールの TensorRT モデルを取得しましたが、このモジュールのモデル入出力は固定サイズであるため、基本的に変換処理に大きな問題はありません。さらに、このモデルの機能は、画像の特徴を取得することであり、時間はかかりますが、一度抽出するだけでよく、その後の入力ポイントやボックス プロンプトに対して繰り返し抽出する必要はありません。フロントエンドとバックエンドの展開を適切に設計できます。

Prompt_Mask モジュールの Pytorch モデルは ONNX モデルに変換されます

上で述べたように、プロンプトエンコーディングとマスクデコーディングモデルは埋め込みに対して操作されており、初期段階で画像埋め込みが抽出された後は、入力プロンプトポイントとボックスの座標を必要に応じて変更するだけです。この部分のモデル変換は公式のスクリプトが提供されており、スクリプトを実行するだけで onnx モデルを取得します。

# clone官方代码
git clone https://github.com/facebookresearch/segment-anything

**注意: **ソース コードのマスクはデコード後の低サイズ マスクであり、入力画像の元のサイズに従って復元する必要がありますが、この元のサイズが onnx 変換中に入力ノードとして使用される場合は、 , TensorRT モデルの変換時にも使用されます. このパラメータを入力し、長さと幅のパラメータ値を固定する必要があります. ただし、ユーザーが入力した画像のサイズは事前に知ることができないため、これはパラメータは個別に抽出する必要がある、つまり、低次元マスクの後処理はモデル パーツとしてではなく個別に処理されるため、ソース コードを少し変更する必要があります。

# 修改"segment_anything/utils/onnx.py"中的"forward"函数为如下：
def forward(
    self,
    image_embeddings: torch.Tensor,
    point_coords: torch.Tensor,
    point_labels: torch.Tensor,
    mask_input: torch.Tensor,
    has_mask_input: torch.Tensor
    # orig_im_size: torch.Tensor,
):
    sparse_embedding = self._embed_points(point_coords, point_labels)
    dense_embedding = self._embed_masks(mask_input, has_mask_input)

    masks, scores = self.model.mask_decoder.predict_masks(
        image_embeddings=image_embeddings,
        image_pe=self.model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embedding,
        dense_prompt_embeddings=dense_embedding,
    )

    if self.use_stability_score:
        scores = calculate_stability_score(
            masks, self.model.mask_threshold, self.stability_score_offset
        )

    if self.return_single_mask:
        masks, scores = self.select_masks(masks, scores, point_coords.shape[1])

    return masks, scores
    # upscaled_masks = self.mask_postprocessing(masks, orig_im_size)

    # if self.return_extra_metrics:
    #     stability_scores = calculate_stability_score(
    #         upscaled_masks, self.model.mask_threshold, self.stability_score_offset
    #     )
    #     areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1)
    #     return upscaled_masks, scores, stability_scores, areas, masks

    # return upscaled_masks, scores, masks

モデル内の上記の関数を変更した後、sam_vit_h_4b8939.pth モデルのプロンプト エンコード部分とマスク デコード部分を onnx モデルに変換します。

# 下载default模型在库下的weights文件夹，并进行onnx模型的转换
python scripts/onnx2trt.py --prompt_masks_pt2onnx 

**注意:** モデル変換プロセス中、opset バージョンは onnx バージョンと一致する必要があります。そうでない場合、TensorRT モデル変換プロセス中にエラーが報告されますが、これは大きな落とし穴です。

Prompt_Mask モジュールの ONNX モデルは TensorRT モデルに変換されます

このリンクでは、入力に含まれるプロンプトポイントとチェンジポイントの正負の数値が可変、つまり入力サイズが動的であるため、変換プロセス中にマルチサイズパラメータを設定する必要があります。以下に続きます:

trtexec --onnx=weights/sam_default_prompt_mask.onnx --workspace=4096 --shapes=image_embeddings:1x256x64x64,point_coords:1x1x2,point_labels:1x1,mask_input:1x1x256x256,has_mask_input:1 --minShapes=image_embeddings:1x256x64x64,point_coords:1x1x2,point_labels:1x1,mask_input:1x1x256x256,has_mask_input:1 --optShapes=image_embeddings:1x256x64x64,point_coords:1x10x2,point_labels:1x10,mask_input:1x1x256x256,has_mask_input:1 --maxShapes=image_embeddings:1x256x64x64,point_coords:1x20x2,point_labels:1x20,mask_input:1x1x256x256,has_mask_input:1 --saveEngine=weights/sam_default_prompt_mask.engine

上記のプロセスを完了すると、TensorRT の 2 つの高速化されたエンジン ファイルが得られ、モデルの推論タスクを実行できるようになります。推論スクリプトを提供します。

python sam_trt_inference.py