TensorRTってなんだ？〜深層学習モデルを爆速化するNVIDIA推論エンジンの全貌〜

# 最適化前（3回のカーネル起動）
Conv → メモリ書き込み → BatchNorm → メモリ書き込み → ReLU

# 最適化後（1回のカーネル起動）
ConvBNReLU（中間結果はGPUレジスタ内で処理）

# CUDA 12.x環境を前提
# TensorRTとONNX関連パッケージのインストール
pip install tensorrt==10.14.* --break-system-packages
pip install onnx==1.18.0 --break-system-packages
pip install onnxruntime-gpu==1.20.0 --break-system-packages
pip install torch torchvision --break-system-packages

python -c "import tensorrt; print(f'TensorRT version: {tensorrt.__version__}')"

# config.yaml - 開発環境用（このままコピーして使える）
# デバッグしやすさを重視した設定

tensorrt:
  precision: "fp16"          # FP16で開発（INT8キャリブレーション不要）
  workspace_size_gb: 4       # ワークスペースサイズ（GB）
  max_batch_size: 1          # デバッグ時は小さいバッチ
  verbose: true              # 詳細ログを出力
  timing_cache: null         # タイミングキャッシュは無効

onnx:
  opset_version: 17          # ONNXオペセットバージョン
  dynamic_axes: true         # 動的な入力サイズを許可
  simplify: true             # ONNXグラフを簡略化

inference:
  device: "cuda:0"           # 使用するGPUデバイス
  warmup_iterations: 5       # ウォームアップ回数
  benchmark_iterations: 100  # ベンチマーク回数

logging:
  level: "DEBUG"             # ログレベル
  file: null                 # ファイル出力は無効

# config.production.yaml - 本番環境用（このままコピーして使える）
# パフォーマンスと安定性を重視した設定

tensorrt:
  precision: "fp16"          # FP16（精度と速度のバランス）
  workspace_size_gb: 8       # ワークスペースサイズ（GB）
  max_batch_size: 32         # 本番向けバッチサイズ
  verbose: false             # ログ抑制
  timing_cache: "/var/cache/tensorrt/timing.cache"
  builder_optimization_level: 5  # 最大最適化

onnx:
  opset_version: 17
  dynamic_axes: false        # 固定サイズで最適化
  simplify: true

inference:
  device: "cuda:0"
  warmup_iterations: 10
  benchmark_iterations: 0    # ベンチマーク無効

logging:
  level: "WARNING"
  file: "/var/log/tensorrt/inference.log"

monitoring:
  enable_metrics: true       # メトリクス収集
  prometheus_port: 9090      # Prometheus用ポート

# config.test.yaml - テスト/CI環境用（このままコピーして使える）
# 再現性と高速なCI実行を重視した設定

tensorrt:
  precision: "fp32"          # FP32で精度検証
  workspace_size_gb: 2       # 最小限のワークスペース
  max_batch_size: 1          # 小さいバッチで高速テスト
  verbose: false
  timing_cache: null
  deterministic: true        # 決定論的な動作を保証

onnx:
  opset_version: 17
  dynamic_axes: false
  simplify: true

inference:
  device: "cuda:0"
  warmup_iterations: 1
  benchmark_iterations: 10

logging:
  level: "ERROR"             # エラーのみ出力
  file: null

validation:
  rtol: 1e-3                 # 相対許容誤差
  atol: 1e-5                 # 絶対許容誤差
  reference_backend: "pytorch"  # 比較対象

"""
TensorRT推論サンプル
PyTorchモデルをONNX経由でTensorRTエンジンに変換し、推論を実行する

使い方: python tensorrt_inference.py
必要なパッケージ: pip install tensorrt onnx torch torchvision
"""

import torch
import torch.onnx
import tensorrt as trt
import numpy as np
from pathlib import Path
import yaml


def load_config(config_path: str = "config.yaml") -> dict:
    """設定ファイルを読み込む"""
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def export_to_onnx(
    model: torch.nn.Module,
    input_shape: tuple,
    onnx_path: str,
    opset_version: int = 17
) -> None:
    """PyTorchモデルをONNXフォーマットでエクスポート"""
    model.eval()
    dummy_input = torch.randn(*input_shape)
    
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=opset_version,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        }
    )
    print(f"ONNXモデルを保存しました: {onnx_path}")


def build_tensorrt_engine(
    onnx_path: str,
    engine_path: str,
    precision: str = "fp16",
    workspace_size_gb: int = 4,
    verbose: bool = False
) -> trt.ICudaEngine:
    """ONNXモデルからTensorRTエンジンをビルド"""
    
    # ロガーの設定
    logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING)
    
    # ビルダーとネットワークの作成
    builder = trt.Builder(logger)
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    
    # ONNXパーサーでモデルを読み込み
    parser = trt.OnnxParser(network, logger)
    with open(onnx_path, "rb") as f:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(f"ONNXパースエラー: {parser.get_error(error)}")
            raise RuntimeError("ONNXモデルのパースに失敗しました")
    
    # ビルダー設定
    config = builder.create_builder_config()
    config.set_memory_pool_limit(
        trt.MemoryPoolType.WORKSPACE,
        workspace_size_gb * (1 << 30)  # GB -> バイト変換
    )
    
    # 精度設定
    if precision == "fp16":
        if builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            print("FP16精度を有効化しました")
        else:
            print("警告: このGPUはFP16に対応していません。FP32で続行します")
    elif precision == "int8":
        if builder.platform_has_fast_int8:
            config.set_flag(trt.BuilderFlag.INT8)
            print("INT8精度を有効化しました（キャリブレーションが必要です）")
        else:
            print("警告: このGPUはINT8に対応していません。FP32で続行します")
    
    # エンジンのビルド
    print("TensorRTエンジンをビルド中...（数分かかる場合があります）")
    serialized_engine = builder.build_serialized_network(network, config)
    
    if serialized_engine is None:
        raise RuntimeError("TensorRTエンジンのビルドに失敗しました")
    
    # エンジンの保存
    with open(engine_path, "wb") as f:
        f.write(serialized_engine)
    print(f"TensorRTエンジンを保存しました: {engine_path}")
    
    # デシリアライズしてエンジンを返す
    runtime = trt.Runtime(logger)
    return runtime.deserialize_cuda_engine(serialized_engine)


def load_tensorrt_engine(engine_path: str) -> trt.ICudaEngine:
    """保存済みTensorRTエンジンを読み込む"""
    logger = trt.Logger(trt.Logger.WARNING)
    runtime = trt.Runtime(logger)
    
    with open(engine_path, "rb") as f:
        return runtime.deserialize_cuda_engine(f.read())


class TensorRTInference:
    """TensorRT推論を実行するクラス"""
    
    def __init__(self, engine: trt.ICudaEngine):
        self.engine = engine
        self.context = engine.create_execution_context()
        
        # 入出力テンソル情報の取得
        self.input_name = engine.get_tensor_name(0)
        self.output_name = engine.get_tensor_name(1)
        
        # CUDAストリームの作成
        import cuda  # cuda-pythonが必要
        self.stream = cuda.cudart.cudaStreamCreate()[1]
    
    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """推論を実行"""
        import cuda
        
        # 入力データの準備
        input_data = np.ascontiguousarray(input_data.astype(np.float32))
        output_shape = self._get_output_shape(input_data.shape)
        output_data = np.empty(output_shape, dtype=np.float32)
        
        # GPUメモリの確保
        d_input = cuda.cudart.cudaMalloc(input_data.nbytes)[1]
        d_output = cuda.cudart.cudaMalloc(output_data.nbytes)[1]
        
        # ホスト→デバイス転送
        cuda.cudart.cudaMemcpyAsync(
            d_input, input_data.ctypes.data,
            input_data.nbytes,
            cuda.cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
            self.stream
        )
        
        # テンソルアドレスの設定
        self.context.set_tensor_address(self.input_name, int(d_input))
        self.context.set_tensor_address(self.output_name, int(d_output))
        
        # 推論実行
        self.context.execute_async_v3(self.stream)
        
        # デバイス→ホスト転送
        cuda.cudart.cudaMemcpyAsync(
            output_data.ctypes.data, d_output,
            output_data.nbytes,
            cuda.cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
            self.stream
        )
        
        # 同期
        cuda.cudart.cudaStreamSynchronize(self.stream)
        
        # メモリ解放
        cuda.cudart.cudaFree(d_input)
        cuda.cudart.cudaFree(d_output)
        
        return output_data
    
    def _get_output_shape(self, input_shape: tuple) -> tuple:
        """出力テンソルのシェイプを取得"""
        self.context.set_input_shape(self.input_name, input_shape)
        return tuple(self.context.get_tensor_shape(self.output_name))


def benchmark(inference: TensorRTInference, input_shape: tuple, iterations: int = 100):
    """推論速度のベンチマーク"""
    import time
    
    # ウォームアップ
    dummy_input = np.random.randn(*input_shape).astype(np.float32)
    for _ in range(10):
        inference.infer(dummy_input)
    
    # ベンチマーク
    times = []
    for _ in range(iterations):
        start = time.perf_counter()
        inference.infer(dummy_input)
        end = time.perf_counter()
        times.append((end - start) * 1000)  # ミリ秒
    
    avg_time = np.mean(times)
    std_time = np.std(times)
    
    print(f"\n===== ベンチマーク結果 =====")
    print(f"平均推論時間: {avg_time:.2f} ms")
    print(f"標準偏差: {std_time:.2f} ms")
    print(f"スループット: {1000 / avg_time:.2f} FPS")
    
    return avg_time


def main():
    """メイン関数"""
    # ResNet18をサンプルモデルとして使用
    from torchvision.models import resnet18, ResNet18_Weights
    
    model = resnet18(weights=ResNet18_Weights.DEFAULT)
    input_shape = (1, 3, 224, 224)
    
    onnx_path = "resnet18.onnx"
    engine_path = "resnet18.engine"
    
    # Step 1: ONNXエクスポート
    if not Path(onnx_path).exists():
        export_to_onnx(model, input_shape, onnx_path)
    
    # Step 2: TensorRTエンジンビルド
    if not Path(engine_path).exists():
        engine = build_tensorrt_engine(
            onnx_path,
            engine_path,
            precision="fp16",
            workspace_size_gb=4,
            verbose=True
        )
    else:
        engine = load_tensorrt_engine(engine_path)
        print(f"既存のTensorRTエンジンを読み込みました: {engine_path}")
    
    # Step 3: 推論実行
    inference = TensorRTInference(engine)
    
    # テスト推論
    test_input = np.random.randn(*input_shape).astype(np.float32)
    output = inference.infer(test_input)
    print(f"\n入力シェイプ: {test_input.shape}")
    print(f"出力シェイプ: {output.shape}")
    
    # ベンチマーク
    benchmark(inference, input_shape, iterations=100)


if __name__ == "__main__":
    main()

$ python tensorrt_inference.py
ONNXモデルを保存しました: resnet18.onnx
FP16精度を有効化しました
TensorRTエンジンをビルド中...（数分かかる場合があります）
TensorRTエンジンを保存しました: resnet18.engine

入力シェイプ: (1, 3, 224, 224)
出力シェイプ: (1, 1000)

===== ベンチマーク結果 =====
平均推論時間: 1.23 ms
標準偏差: 0.15 ms
スループット: 813.01 FPS

"""
リアルタイム画像分類向けTensorRT推論
レイテンシ重視の設定でストリーミング処理を実現
"""

import cv2
import numpy as np
import tensorrt as trt
from typing import Tuple


class RealtimeClassifier:
    """リアルタイム画像分類クラス"""
    
    def __init__(self, engine_path: str, labels_path: str):
        self.engine = self._load_engine(engine_path)
        self.context = self.engine.create_execution_context()
        self.labels = self._load_labels(labels_path)
        
        # 推論用バッファの事前確保（レイテンシ削減）
        self._allocate_buffers()
    
    def _load_engine(self, path: str) -> trt.ICudaEngine:
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)
        with open(path, "rb") as f:
            return runtime.deserialize_cuda_engine(f.read())
    
    def _load_labels(self, path: str) -> list:
        with open(path, "r") as f:
            return [line.strip() for line in f]
    
    def _allocate_buffers(self):
        """CUDAバッファの事前確保"""
        import cuda
        
        self.input_shape = (1, 3, 224, 224)
        self.output_shape = (1, 1000)
        
        input_size = np.prod(self.input_shape) * 4  # float32
        output_size = np.prod(self.output_shape) * 4
        
        self.d_input = cuda.cudart.cudaMalloc(input_size)[1]
        self.d_output = cuda.cudart.cudaMalloc(output_size)[1]
        self.h_output = np.empty(self.output_shape, dtype=np.float32)
        self.stream = cuda.cudart.cudaStreamCreate()[1]
    
    def preprocess(self, frame: np.ndarray) -> np.ndarray:
        """画像の前処理"""
        # リサイズ
        img = cv2.resize(frame, (224, 224))
        # BGR -> RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # 正規化
        img = img.astype(np.float32) / 255.0
        img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
        # NHWC -> NCHW
        img = np.transpose(img, (2, 0, 1))
        return np.expand_dims(img, axis=0).astype(np.float32)
    
    def classify(self, frame: np.ndarray) -> Tuple[str, float]:
        """画像を分類し、クラス名と確信度を返す"""
        import cuda
        
        # 前処理
        input_data = np.ascontiguousarray(self.preprocess(frame))
        
        # 推論（事前確保バッファを使用）
        cuda.cudart.cudaMemcpyAsync(
            self.d_input, input_data.ctypes.data,
            input_data.nbytes,
            cuda.cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
            self.stream
        )
        
        input_name = self.engine.get_tensor_name(0)
        output_name = self.engine.get_tensor_name(1)
        self.context.set_tensor_address(input_name, int(self.d_input))
        self.context.set_tensor_address(output_name, int(self.d_output))
        
        self.context.execute_async_v3(self.stream)
        
        cuda.cudart.cudaMemcpyAsync(
            self.h_output.ctypes.data, self.d_output,
            self.h_output.nbytes,
            cuda.cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
            self.stream
        )
        cuda.cudart.cudaStreamSynchronize(self.stream)
        
        # 結果の解釈
        probs = np.exp(self.h_output) / np.sum(np.exp(self.h_output))
        class_idx = np.argmax(probs)
        confidence = probs[0, class_idx]
        
        return self.labels[class_idx], confidence


def main():
    classifier = RealtimeClassifier(
        engine_path="resnet18.engine",
        labels_path="imagenet_labels.txt"
    )
    
    cap = cv2.VideoCapture(0)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        label, confidence = classifier.classify(frame)
        
        # 結果を表示
        text = f"{label}: {confidence:.2%}"
        cv2.putText(frame, text, (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.imshow("Realtime Classification", frame)
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()

"""
バッチ処理向けTensorRT推論
スループット重視で大量データを効率的に処理
"""

import numpy as np
import tensorrt as trt
from pathlib import Path
from typing import Generator, List
import time


class BatchInference:
    """バッチ推論クラス"""
    
    def __init__(self, engine_path: str, batch_size: int = 32):
        self.batch_size = batch_size
        self.engine = self._load_engine(engine_path)
        self.context = self.engine.create_execution_context()
    
    def _load_engine(self, path: str) -> trt.ICudaEngine:
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)
        with open(path, "rb") as f:
            return runtime.deserialize_cuda_engine(f.read())
    
    def infer_batch(self, images: np.ndarray) -> np.ndarray:
        """バッチ推論を実行"""
        import cuda
        
        actual_batch = images.shape[0]
        input_name = self.engine.get_tensor_name(0)
        output_name = self.engine.get_tensor_name(1)
        
        # 入力シェイプを動的に設定
        self.context.set_input_shape(input_name, images.shape)
        output_shape = tuple(self.context.get_tensor_shape(output_name))
        
        # メモリ確保
        images = np.ascontiguousarray(images.astype(np.float32))
        output = np.empty(output_shape, dtype=np.float32)
        
        d_input = cuda.cudart.cudaMalloc(images.nbytes)[1]
        d_output = cuda.cudart.cudaMalloc(output.nbytes)[1]
        stream = cuda.cudart.cudaStreamCreate()[1]
        
        # 転送と推論
        cuda.cudart.cudaMemcpyAsync(
            d_input, images.ctypes.data, images.nbytes,
            cuda.cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream
        )
        
        self.context.set_tensor_address(input_name, int(d_input))
        self.context.set_tensor_address(output_name, int(d_output))
        self.context.execute_async_v3(stream)
        
        cuda.cudart.cudaMemcpyAsync(
            output.ctypes.data, d_output, output.nbytes,
            cuda.cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream
        )
        cuda.cudart.cudaStreamSynchronize(stream)
        
        # クリーンアップ
        cuda.cudart.cudaFree(d_input)
        cuda.cudart.cudaFree(d_output)
        cuda.cudart.cudaStreamDestroy(stream)
        
        return output
    
    def process_dataset(
        self, 
        image_paths: List[str],
        preprocess_fn
    ) -> Generator[np.ndarray, None, None]:
        """データセット全体をバッチ処理"""
        total = len(image_paths)
        processed = 0
        
        for i in range(0, total, self.batch_size):
            batch_paths = image_paths[i:i + self.batch_size]
            batch_images = np.stack([
                preprocess_fn(path) for path in batch_paths
            ])
            
            yield self.infer_batch(batch_images)
            
            processed += len(batch_paths)
            print(f"処理中: {processed}/{total} ({processed/total:.1%})")


def preprocess_image(path: str) -> np.ndarray:
    """画像の前処理"""
    import cv2
    img = cv2.imread(path)
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
    return np.transpose(img, (2, 0, 1))


def main():
    # 画像パスのリスト（例）
    image_dir = Path("./images")
    image_paths = list(image_dir.glob("*.jpg"))
    
    print(f"処理対象: {len(image_paths)} 画像")
    
    # バッチ推論
    batch_inference = BatchInference(
        engine_path="resnet18.engine",
        batch_size=32
    )
    
    start = time.time()
    results = []
    for batch_output in batch_inference.process_dataset(
        [str(p) for p in image_paths],
        preprocess_image
    ):
        results.append(batch_output)
    
    elapsed = time.time() - start
    throughput = len(image_paths) / elapsed
    
    print(f"\n===== 処理完了 =====")
    print(f"処理時間: {elapsed:.2f} 秒")
    print(f"スループット: {throughput:.2f} 画像/秒")


if __name__ == "__main__":
    main()

"""
Torch-TensorRTによる簡易最適化
既存のPyTorchコードを最小限の変更で高速化
"""

import torch
import torch_tensorrt
from torchvision.models import resnet18, ResNet18_Weights
import time


def optimize_with_torch_tensorrt():
    """Torch-TensorRTでモデルを最適化"""
    
    # PyTorchモデルのロード
    model = resnet18(weights=ResNet18_Weights.DEFAULT).eval().cuda()
    
    # サンプル入力
    sample_input = torch.randn(1, 3, 224, 224).cuda()
    
    # Torch-TensorRTによる最適化（ワンライナー）
    optimized_model = torch_tensorrt.compile(
        model,
        inputs=[
            torch_tensorrt.Input(
                shape=[1, 3, 224, 224],
                dtype=torch.float32
            )
        ],
        enabled_precisions={torch.float16},  # FP16を有効化
        workspace_size=1 << 30,  # 1GB
        truncate_double=True
    )
    
    return model, optimized_model


def compare_performance(original, optimized, iterations=100):
    """オリジナルと最適化版のパフォーマンス比較"""
    
    sample_input = torch.randn(1, 3, 224, 224).cuda()
    
    # ウォームアップ
    for _ in range(10):
        _ = original(sample_input)
        _ = optimized(sample_input)
    
    torch.cuda.synchronize()
    
    # オリジナルのベンチマーク
    start = time.perf_counter()
    for _ in range(iterations):
        _ = original(sample_input)
        torch.cuda.synchronize()
    original_time = (time.perf_counter() - start) / iterations * 1000
    
    # 最適化版のベンチマーク
    start = time.perf_counter()
    for _ in range(iterations):
        _ = optimized(sample_input)
        torch.cuda.synchronize()
    optimized_time = (time.perf_counter() - start) / iterations * 1000
    
    speedup = original_time / optimized_time
    
    print(f"\n===== パフォーマンス比較 =====")
    print(f"PyTorch (オリジナル): {original_time:.2f} ms")
    print(f"Torch-TensorRT: {optimized_time:.2f} ms")
    print(f"高速化率: {speedup:.2f}x")
    
    return speedup


def main():
    print("Torch-TensorRTでモデルを最適化中...")
    original, optimized = optimize_with_torch_tensorrt()
    
    print("パフォーマンス比較を実行中...")
    compare_performance(original, optimized)
    
    # 最適化モデルの保存
    torch.jit.save(
        torch.jit.trace(optimized, torch.randn(1, 3, 224, 224).cuda()),
        "resnet18_trt.ts"
    )
    print("\n最適化モデルを保存しました: resnet18_trt.ts")


if __name__ == "__main__":
    main()