基本的な使い方

AcuiRTによるモデル変換は以下の3ステップから構成されます。

1. モデル・データセットを作成する

高速化対象のモデルとデータセットを作成します。 モデルはPyTorchのnn.Module形式で作成してください。

2. 変換方法を指定する

変換方法をconfigを記述して指定します。AcuiRTでは、現在下記の深層学習コンパイラをサポートしています。

• TensorRT

3. 変換して最適化されたモデルを実行する

Python コードを実行して変換を行います。以下はResNet50に対してPTQによるint8量子化を適用してTensorRTに変換する例です。

import torch
from acuirt.convert.convert import convert_model
from acuirt.inference.inference import load_runtime_modules
from torchvision.models import resnet50


def main():
    resnet = resnet50()
    resnet = resnet.cuda().eval()

    # int8量子化(PTQ)でTensorRTに変換する場合の設定
    config = {
        "rt_mode": "onnx",
        "auto": True,
        "int8": True,
    }

    # 変換後のモデルを保存するパスを指定してください。
    path = "/path/to/save/model"

    # ダミーのデータセットを作成します
    # Iterableなデータセットを渡すことで、calibrationが自動的に行われます。
    data = [((torch.randn(1, 3, 224, 224), ), {}) for _ in range(10)]

    # TensorRTへの変換と、calibrationを実行します。
    # 変換後のモデルはpathに保存されます。
    # また、summaryというdict型の変数に変換されたモデルの情報が格納されます。
    summary = convert_model(resnet, config, path, False, data)

    # TensorRT用の推論エンジンをロードします。
    model = load_runtime_modules(resnet, summary, path)

    # 推論を実行します。
    args, _ = data[0]
    args = [arg.cuda() for arg in args]
    model(*args)


if __name__ == "__main__":
    main()