環境

Jetson Orin Nano developper kit 8GB
※下記のJetpackからアップデートを試みましたが失敗したため、SDイメージ書き直しを実施しました。
Jetpack 6.1（nvidia-l4t-ore 36.4.0-20240912212859）

Super化手順

下記のサイトにアクセスしてSDイメージをダウンロード
https://developer.nvidia.com/embedded/jetpack
※下記のようにJetpack6を使っている人は下記のコマンドでMAXNモードになると書いてありましたが、自分の場合はMAXNモードが表示されませんでした。

　
balenaEtcher 等を利用してSDカードにイメージを書き込む。
　
初期設定とログイン後、一度rebootが必要です。
再起動後にnvidia-l4t-core 36.4.2になりました。
　
下記のコマンドを実行し、MAXNモードを設定。
```
sudo rm -rf /etc/nvpmodel.conf
```
Power modeを表示してMAXNモードを確認、SUPER！

試しにmovileViTの速度比較(torch, onnx, tensorrt)

追記@2025/1/12

環境構築

追記@2025/1/13

cuSPARSELtのインストール(pytorchに必要です）

wget https://developer.download.nvidia.com/compute/cusparselt/0.6.3/local_installers/cusparselt-local-tegra-repo-ubuntu2204-0.6.3_1.0-1_arm64.debsudo dpkg -i cusparselt-local-tegra-repo-ubuntu2204-0.6.3_1.0-1_arm64.debsudo cp /var/cusparselt-local-tegra-repo-ubuntu2204-0.6.3/cusparselt-*-keyring.gpg /usr/share/keyrings/sudo apt-get updatesudo apt-get -y install libcusparselt0 libcusparselt-dev

pytorch(cuda有）とtorchvisionの準備

pip install http://jetson.webredirect.org/jp6/cu126/+f/5cf/9ed17e35cb752/torch-2.5.0-cp310-cp310-linux_aarch64.whl#sha256=5cf9ed17e35cb7523812aeda9e7d6353c437048c5a6df1dc6617650333049092

pip install http://jetson.webredirect.org/jp6/cu126/+f/5f9/67f920de3953f/torchvision-0.20.0-cp310-cp310-linux_aarch64.whl#sha256=5f967f920de3953f2a39d95154b1feffd5ccc06b4589e51540dc070021a9adb9

参考：https://forums.developer.nvidia.com/t/pytorch-and-torvision-version-issue-runtimeerror-operator-torchvision-nms-does-not-exist/312446/4

pycuda
環境変数の設定→.barshrcに書く

export CPATH=/usr/local/cuda/include:$CPATH
export LIBRARY_PATH=/usr/local/cuda/lib64:$LIBRARY_PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export PATH=/usr/local/cuda/bin:$PATH
pip install pycuda

再インストール後にPythonを起動し、以下のコードを実行して確認。

import torch
print("PyTorch version:", torch.__version__)
print("Compiled CUDA version:", torch.version.cuda)
print("Is CUDA available?:", torch.cuda.is_available())

試験条件

推論モデル：mobilevit_xxs
　https://huggingface.co/timm/mobilevit_xxs.cvnets_in1k
ダミー入力：torch.randn(1, 3, 224, 224)）
ウォームアップ5回のち、１０回推論を実施
平均推論時間(ms)を下記のテーブルに表示

結果

MODE	torch	onnx	tensorrt(execute_v2)
7W	50.00	※落ち	-
15W	26.73	26.76	4.16
Super	23.75	23.63	4.85

※onnxの最適化でコアダンプの模様

コメント

tensorrt早すぎるﾅﾆｺﾚ、Superにしたら逆転...
　（ちなみに推論開始までは１分くらい、一番待った）
-Superにするとtorchもonnxも15Wから約３ms（1割強）ずつ早くなっているが、
　+10Wにしては微妙な結果？
今回はonnxの最適化も試してみましたが、なぜか7Wモードでは落ちてしまいました。
まだよくわからないことが多いので継続調査します

サンプルコード

import torch
from torch import nn
import timm
import time
import onnx
import onnxruntime as ort
import numpy as np
import os

# TensorRT 関連のライブラリがあるかチェック
try:
    import tensorrt as trt
    import pycuda.driver as cuda
    import pycuda.autoinit
    TENSORRT_AVAILABLE = True
except ImportError:
    print("Warning: TensorRT or PyCUDA が見つかりませんでした。TensorRT 推論をスキップします。")
    TENSORRT_AVAILABLE = False

class MobileViTModel(nn.Module):
    """MobileViT を用いた単一カメラモデル"""
    def __init__(self, output_dim=2, pretrained=True, model_name='mobilevit_xxs'):
        super(MobileViTModel, self).__init__()
        self.mobilevit = timm.create_model(model_name, pretrained=pretrained)
        num_ftrs = self.mobilevit.get_classifier().in_features
        self.mobilevit.reset_classifier(0)  # 分類器をリセット
        self.fc = nn.Linear(num_ftrs, output_dim)
        self.tanh = nn.Tanh()

    def forward(self, x):
        x = self.mobilevit(x)
        x = self.fc(x)
        x = self.tanh(x)  # 出力を -1 ~ 1 に制限
        return x

def measure_pytorch_inference(model, input_tensor, device, warm_up=5, runs=10):
    """通常のPyTorchモデルの推論時間を測定（ウォームアップランと複数回実行を含む）"""
    model.to(device)
    input_tensor = input_tensor.to(device)
    model.eval()
    timings = []
    with torch.no_grad():
        # ウォームアップラン
        for _ in range(warm_up):
            _ = model(input_tensor)
        # 推論時間の測定
        for _ in range(runs):
            if device.type == 'cuda':
                torch.cuda.synchronize()
            start_time = time.time()
            _ = model(input_tensor)
            if device.type == 'cuda':
                torch.cuda.synchronize()
            end_time = time.time()
            inference_time = (end_time - start_time) * 1000  # ミリ秒
            timings.append(inference_time)
    avg_time = sum(timings) / len(timings)
    std_time = (sum((x - avg_time) ** 2 for x in timings) / len(timings)) ** 0.5
    return avg_time, std_time

def measure_onnxruntime_inference(onnx_model_path, input_tensor, warm_up=5, runs=10):
    """ONNX Runtimeモデルの推論時間を測定（ウォームアップランと複数回実行を含む）"""
    # セッションオプションの設定
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED  # 高度な最適化
    sess_options.intra_op_num_threads = 4  # スレッド数を設定（必要に応じて調整）
    sess_options.inter_op_num_threads = 4

    ort_session = ort.InferenceSession(onnx_model_path, sess_options=sess_options)
    ort_inputs = {ort_session.get_inputs()[0].name: input_tensor.cpu().numpy()}
    timings = []
    # ウォームアップラン
    for _ in range(warm_up):
        _ = ort_session.run(None, ort_inputs)
    # 推論時間の測定
    for _ in range(runs):
        start_time = time.time()
        _ = ort_session.run(None, ort_inputs)
        end_time = time.time()
        inference_time = (end_time - start_time) * 1000  # ミリ秒
        timings.append(inference_time)
    avg_time = sum(timings) / len(timings)
    std_time = (sum((x - avg_time) ** 2 for x in timings) / len(timings)) ** 0.5
    return avg_time, std_time

# -------------------------------------------------------------------
# 以下、TensorRT 8.5+ の API に合わせたビルド＆推論関数を定義
# -------------------------------------------------------------------

def build_tensorrt_engine(onnx_file_path, max_workspace_size=(1 << 30)):
    """
    TensorRT 8.5+ でのエンジンビルド:
      1) builder.build_serialized_network(network, config) でシリアライズ
      2) runtime.deserialize_cuda_engine(serialized_engine) でエンジン生成
    """
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(TRT_LOGGER)

    # EXPLICIT_BATCH: 動的形状用
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    parser = trt.OnnxParser(network, TRT_LOGGER)

    config = builder.create_builder_config()
    # TensorRT 8.4 以降は config.set_memory_pool_limit でワークスペース設定
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, max_workspace_size)

    # ONNX をパース
    with open(onnx_file_path, "rb") as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(parser.get_error(i))
            raise RuntimeError("Failed to parse the ONNX file.")

    # 例: バッチサイズ1, 3ch, 224x224 の固定入力にする
    profile = builder.create_optimization_profile()
    input_name = network.get_input(0).name
    profile.set_shape(
        input_name,
        (1, 3, 224, 224),  # 最小形状
        (1, 3, 224, 224),  # 最適形状
        (1, 3, 224, 224),  # 最大形状
    )
    config.add_optimization_profile(profile)

    # シリアライズエンジンをビルド
    serialized_engine = builder.build_serialized_network(network, config)
    if serialized_engine is None:
        raise RuntimeError("Failed to build serialized engine.")

    # Runtime を作ってデシリアライズ
    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(serialized_engine)
    if engine is None:
        raise RuntimeError("Failed to deserialize engine.")

    return engine

def measure_tensorrt_inference(onnx_model_path, input_tensor, warm_up=5, runs=10):
    """
    TensorRT 10.3+ の execute_async_v3() を使用して推論時間を測定。
    """
    engine = build_tensorrt_engine(onnx_model_path)
    context = engine.create_execution_context()

    # 入力と出力のテンソル名を取得
    input_name, output_name = None, None
    for i in range(engine.num_io_tensors):
        tname = engine.get_tensor_name(i)
        mode = engine.get_tensor_mode(tname)
        if mode == trt.TensorIOMode.INPUT:
            input_name = tname
        elif mode == trt.TensorIOMode.OUTPUT:
            output_name = tname

    if not input_name or not output_name:
        raise RuntimeError("入力または出力のテンソル名が見つかりません。")

    # 入力の形状を設定
    shape = tuple(input_tensor.shape)
    context.set_input_shape(input_name, shape)

    # ホスト側のデータを用意
    host_input = input_tensor.cpu().numpy().astype(np.float32)
    input_size = host_input.nbytes

    # 出力データの形状を仮定（例: (batch, 2)）
    output_shape = (shape[0], 2)
    host_output = np.empty(output_shape, dtype=np.float32)
    output_size = host_output.nbytes

    # デバイスメモリの確保
    d_input = cuda.mem_alloc(input_size)
    d_output = cuda.mem_alloc(output_size)

    # テンソルにデバイスメモリを設定
    context.set_tensor_address(input_name, int(d_input))
    context.set_tensor_address(output_name, int(d_output))

    # CUDA ストリームの作成
    stream = cuda.Stream()
    stream_handle = stream.handle

    # ウォームアップ + 計測
    timings = []
    for i in range(warm_up + runs):
        # ホストからデバイスへデータ転送
        cuda.memcpy_htod_async(d_input, host_input, stream)

        if i >= warm_up:
            start = time.time()

        # 推論実行（非同期）
        success = context.execute_async_v3(stream_handle)
        if not success:
            raise RuntimeError("TensorRT 推論実行に失敗しました。")

        # デバイスからホストへデータ転送
        cuda.memcpy_dtoh_async(host_output, d_output, stream)

        # ストリームの同期
        stream.synchronize()

        if i >= warm_up:
            end = time.time()
            inference_time = (end - start) * 1000  # ミリ秒
            timings.append(inference_time)

    # メモリの解放
    d_input.free()
    d_output.free()

    # 統計の計算
    avg_time = sum(timings) / len(timings)
    std_time = (sum((t - avg_time) ** 2 for t in timings) / len(timings)) ** 0.5
    return avg_time, std_time

def measure_tensorrt_inference_v2(onnx_model_path, input_tensor, warm_up=5, runs=10):
    """
    TensorRT 10.3+ の execute_v2() を使用して推論時間を測定。
    
    Parameters:
        onnx_model_path (str): ONNX モデルファイルのパス。
        input_tensor (torch.Tensor): 推論に使用する入力テンソル。
        warm_up (int): ウォームアップランの回数。
        runs (int): 推論を実行する回数。
    
    Returns:
        avg_time (float): 平均推論時間（ミリ秒）。
        std_time (float): 推論時間の標準偏差（ミリ秒）。
    """
    # TensorRT エンジンのビルド
    engine = build_tensorrt_engine(onnx_model_path)
    context = engine.create_execution_context()

    # 入力と出力のテンソル名を取得
    input_name, output_name = None, None
    for i in range(engine.num_io_tensors):
        tname = engine.get_tensor_name(i)
        mode = engine.get_tensor_mode(tname)  # INPUT or OUTPUT
        if mode == trt.TensorIOMode.INPUT:
            input_name = tname
        elif mode == trt.TensorIOMode.OUTPUT:
            output_name = tname

    if not input_name or not output_name:
        raise RuntimeError("入力または出力のテンソル名が見つかりません。")

    # 入力の形状を設定（動的形状の場合のみ）
    shape = tuple(input_tensor.shape)
    if hasattr(context, 'set_input_shape'):
        context.set_input_shape(input_name, shape)

    # ホスト側のデータを用意
    host_input = input_tensor.cpu().numpy().astype(np.float32)
    input_size = host_input.nbytes

    # 出力データの形状を仮定（例: (batch, 2)）
    output_shape = (shape[0], 2)
    host_output = np.empty(output_shape, dtype=np.float32)
    output_size = host_output.nbytes

    # デバイスメモリの確保
    d_input = cuda.mem_alloc(input_size)
    d_output = cuda.mem_alloc(output_size)

    # テンソル名とデバイスメモリのアドレスを設定
    context.set_tensor_address(input_name, int(d_input))
    context.set_tensor_address(output_name, int(d_output))

    # バインディングアドレスをテンソル順にリスト化
    bindings = []
    for i in range(engine.num_io_tensors):
        tname = engine.get_tensor_name(i)
        if tname == input_name:
            bindings.append(int(d_input))
        elif tname == output_name:
            bindings.append(int(d_output))
        else:
            bindings.append(0)  # 他のテンソルがある場合は0を設定（必要に応じて調整）

    # ウォームアップ + 計測
    timings = []
    for i in range(warm_up + runs):
        # ホストからデバイスへデータ転送
        cuda.memcpy_htod(d_input, host_input)

        if i >= warm_up:
            start_time = time.time()

        # 推論実行 (同期)
        success = context.execute_v2(bindings=bindings)
        if not success:
            raise RuntimeError("TensorRT 推論実行に失敗しました。")

        if i >= warm_up:
            end_time = time.time()
            inference_time = (end_time - start_time) * 1000  # ミリ秒
            timings.append(inference_time)

        # デバイスからホストへデータ転送
        cuda.memcpy_dtoh(host_output, d_output)

    # メモリの解放
    d_input.free()
    d_output.free()

    # 統計の計算
    avg_time = sum(timings) / len(timings)
    std_time = (sum((t - avg_time) ** 2 for t in timings) / len(timings)) ** 0.5
    return avg_time, std_time

def measure_tensorrt_inference_async_v3(onnx_model_path, input_tensor, warm_up=5, runs=10):
    """
    TensorRT 10.3+ の execute_async_v3() を使用して推論時間を測定。
    """
    engine = build_tensorrt_engine(onnx_model_path)
    context = engine.create_execution_context()

    # 入力と出力のテンソル名を取得
    input_name, output_name = None, None
    for i in range(engine.num_io_tensors):
        tname = engine.get_tensor_name(i)
        mode = engine.get_tensor_mode(tname)  # INPUT or OUTPUT
        if mode == trt.TensorIOMode.INPUT:
            input_name = tname
        elif mode == trt.TensorIOMode.OUTPUT:
            output_name = tname

    if not input_name or not output_name:
        raise RuntimeError("入力または出力のテンソル名が見つかりません。")

    # 入力の形状を設定
    shape = tuple(input_tensor.shape)
    context.set_input_shape(input_name, shape)

    # ホスト側のデータを用意
    host_input = input_tensor.cpu().numpy().astype(np.float32)
    input_size = host_input.nbytes

    # 出力データの形状を仮定（例: (batch, 2)）
    output_shape = (shape[0], 2)
    host_output = np.empty(output_shape, dtype=np.float32)
    output_size = host_output.nbytes

    # デバイスメモリの確保
    d_input = cuda.mem_alloc(input_size)
    d_output = cuda.mem_alloc(output_size)

    # テンソル名とデバイスメモリのアドレスを設定
    context.set_tensor_address(input_name, int(d_input))
    context.set_tensor_address(output_name, int(d_output))

    # CUDA ストリームの作成
    stream = cuda.Stream()
    stream_handle = stream.handle

    # ウォームアップ + 計測
    timings = []
    for i in range(warm_up + runs):
        # ホストからデバイスへデータ転送 (非同期)
        cuda.memcpy_htod_async(d_input, host_input, stream)

        if i >= warm_up:
            start_time = time.time()

        # 推論実行 (非同期)
        success = context.execute_async_v3(stream_handle)
        if not success:
            raise RuntimeError("TensorRT 推論実行に失敗しました。")

        # デバイスからホストへデータ転送 (非同期)
        cuda.memcpy_dtoh_async(host_output, d_output, stream)

        # ストリームの同期
        stream.synchronize()

        if i >= warm_up:
            end_time = time.time()
            inference_time = (end_time - start_time) * 1000  # ミリ秒
            timings.append(inference_time)

    # メモリの解放
    d_input.free()
    d_output.free()

    # 統計の計算
    avg_time = sum(timings) / len(timings)
    std_time = (sum((t - avg_time) ** 2 for t in timings) / len(timings)) ** 0.5
    return avg_time, std_time, host_output

# -------------------------------------------------------------------
# メイン関数で PyTorch / ONNX Runtime / TensorRT の推論時間をまとめて測定
# -------------------------------------------------------------------
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用デバイス: {device}")

    # モデルのインスタンス化
    model = MobileViTModel(output_dim=2, pretrained=True, model_name='mobilevit_xxs')

    # ダミー入力の作成 (バッチサイズ1、チャンネル3、224x224)
    dummy_input = torch.randn(1, 3, 224, 224)

    # 推論時間を格納する辞書
    results = {}

    # 1. PyTorch モデルの推論時間測定
    print("測定中: PyTorch モデル")
    pytorch_time, pytorch_std = measure_pytorch_inference(model, dummy_input, device, warm_up=5, runs=10)
    results["PyTorch"] = {"平均": pytorch_time, "標準偏差": pytorch_std}
    print(f"PyTorch 推論時間 (平均): {pytorch_time:.2f} ms, 標準偏差: {pytorch_std:.2f} ms\n")

    # 2. ONNX Runtime の推論時間測定
    print("測定中: ONNX Runtime モデル")
    onnx_model_path = "mobilevit_model.onnx"

    try:
        # （エクスポート時は CPU モードを推奨）
        model.cpu()
        cpu_input = dummy_input.cpu()

        # モデルを ONNX 形式でエクスポート
        torch.onnx.export(
            model,
            cpu_input,
            onnx_model_path,
            export_params=True,
            opset_version=14,
            do_constant_folding=True,
            input_names=["input"],
            output_names=["output"],
            dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        )
        print(f"ONNX形式にエクスポート完了: {onnx_model_path} (opset_version=14)")

        # ONNX Runtime での推論
        onnx_time, onnx_std = measure_onnxruntime_inference(onnx_model_path, cpu_input, warm_up=5, runs=10)
        results["ONNX Runtime"] = {"平均": onnx_time, "標準偏差": onnx_std}
        print(f"ONNX Runtime 推論時間 (平均): {onnx_time:.2f} ms, 標準偏差: {onnx_std:.2f} ms\n")

    except Exception as e:
        print(f"ONNX Runtime での推論に失敗: {e}")
        print("ONNX Runtime の推論時間測定をスキップします。\n")

    # 3. TensorRT (10.3+) の推論時間測定
    if TENSORRT_AVAILABLE:
        # a. execute_v2() の推論時間測定
        print("測定中: TensorRT モデル (execute_v2())")
        try:
            trt_time_v2, trt_std_v2 = measure_tensorrt_inference_v2(onnx_model_path, cpu_input, warm_up=5, runs=10)
            results["TensorRT (execute_v2)"] = {"平均": trt_time_v2, "標準偏差": trt_std_v2}
            print(f"TensorRT execute_v2() 推論時間 (平均): {trt_time_v2:.2f} ms, 標準偏差: {trt_std_v2:.2f} ms\n")
        except Exception as e:
            print(f"TensorRT execute_v2() での推論に失敗しました: {e}")
            print("TensorRT execute_v2() の推論時間測定をスキップします。\n")

        # b. execute_async_v3() の推論時間測定
        print("測定中: TensorRT モデル (execute_async_v3())")
        try:
            trt_time_async, trt_std_async, trt_output = measure_tensorrt_inference_async_v3(
                onnx_model_path, cpu_input, warm_up=5, runs=10
            )
            results["TensorRT (execute_async_v3)"] = {"平均": trt_time_async, "標準偏差": trt_std_async}
            print(f"TensorRT execute_async_v3() 推論時間 (平均): {trt_time_async:.2f} ms, 標準偏差: {trt_std_async:.2f} ms")
            print("TensorRT execute_async_v3() 出力データ形状:", trt_output.shape)
            print("TensorRT execute_async_v3() 出力データ (第一行):", trt_output[0] if trt_output.size > 0 else "No data")
            results["TensorRT Output"] = trt_output
            print()
        except Exception as e:
            print(f"TensorRT execute_async_v3() での推論に失敗しました: {e}")
            print("TensorRT execute_async_v3() の推論時間測定をスキップします。\n")
    else:
        print("TensorRT が利用できないためスキップします。\n")

    # 結果の比較表示
    print("### 推論時間の比較 ###")
    for key, value in results.items():
        if key == "TensorRT Output":
            continue  # 出力データはここでは表示しない
        print(f"{key}: 平均 {value['平均']:.2f} ms, 標準偏差: {value['標準偏差']:.2f} ms")

    # オプション: TensorRT 出力と他のフレームワークの出力の比較
    if "TensorRT Output" in results and "PyTorch" in results:
        print("\n### TensorRT 出力と PyTorch 出力の比較 ###")
        # TensorRT 出力
        trt_output = results["TensorRT Output"]
        # PyTorch 出力（実行後に保存されていないので再度実行）
        model.to("cpu")
        model.eval()
        with torch.no_grad():
            pytorch_output = model(dummy_input.cpu()).numpy()
        # 出力の差を計算
        difference = np.abs(trt_output - pytorch_output)
        print("出力データ形状:", trt_output.shape)
        print("出力データの最大差異:", np.max(difference))
        print("出力データの平均差異:", np.mean(difference))
    else:
        print("\nTensorRT 出力の比較は行いませんでした。")

if __name__ == "__main__":
    main()