LoginSignup
1
2

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@Romihi(Romi Hi)

Jetson Orin Nano をSuperにする、試しにmovileViTの速度比較(torch, onnx, tensorrt)

Last updated at Posted at 2025-01-11

環境

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

Super化手順

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

  2. balenaEtcher 等を利用してSDカードにイメージを書き込む。
    　

  3. 初期設定とログイン後、一度rebootが必要です。
    再起動後にnvidia-l4t-core 36.4.2になりました。
    　

  4. 下記のコマンドを実行し、MAXNモードを設定。

    sudo rm -rf /etc/nvpmodel.conf

  5. Power modeを表示してMAXNモードを確認、SUPER！
    image.png

試しに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())

試験条件

結果

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()
1
2
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
1
2

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?