110msから28msへ：リアルタイム推論のためのPyTorchモデル最適化深掘り

pip install torch torchvision onnx onnxruntime fastapi "uvicorn[standard]" numpy psutil

# baseline_server.py
from fastapi import FastAPI
import torch
import torchvision.models as models
import time
import numpy as np

app = FastAPI()

# 1. 事前学習済みFP32モデルを読み込む
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.eval() # 評価モードに設定

# 推論用の勾配計算を無効にするコンテキストマネージャを使用
@torch.inference_mode()
@app.post("/predict")
async def predict():
    # 2. ダミー入力テンソルを作成（バッチサイズ1）
    # これはユーザーからの单个のビデオフレームをシミュレート
    dummy_input = torch.randn(1, 3, 224, 224)

    # 3. 推論時間をベンチマーク
    start_time = time.perf_counter()
    output = model(dummy_input)
    end_time = time.perf_counter()

    latency_ms = (end_time - start_time) * 1000
    
    # 4. 結果を返す
    # 実際のアプリでは、これはポーズデータになります（例：関節座標）
    probabilities = torch.nn.functional.softmax(output, dim=0)
    top5_prob, top5_catid = torch.topk(probabilities, 5)

    return {
        "latency_ms": f"{latency_ms:.2f}",
        "top_category": int(top5_catid),
        "top_probability": float(top5_prob)
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

curl -X POST http://127.0.0.1:8000/predict

{"latency_ms":"110.54","top_category":714,"top_probability":0.0069...}

# quantize_model.py
import torch
import torchvision.models as models
from torch.quantization import quantize_dynamic, get_default_qconfig, prepare, convert

def calibrate_model(model, data_loader, num_batches=10):
    model.eval()
    with torch.no_grad():
        for i, (images, _) in enumerate(data_loader):
            if i >= num_batches:
                break
            model(images)

def main():
    # 1. 元のFP32モデルを読み込む
    float_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
    float_model.eval()

    # 2. 量子化設定を指定
    # 'fbgemm'はx86 CPU recommended バックエンド
    qconfig = get_default_qconfig('fbgemm')
    float_model.qconfig = qconfig

    # 3. 静的量子化用にモデルを準備
    # これは統計を収集する'observer'モジュールを挿入
    prepared_model = prepare(float_model, inplace=False)

    # 4. サンプルデータで較正
    # 実際のシナリオでは代表的な検証データセットを使用
    # ここではデモンストレーション用にランダムデータを使用
    print("Calibrating the model...")
    dummy_dataset = [(torch.randn(3, 224, 224), 0) for _ in range(100)]
    dummy_loader = torch.utils.data.DataLoader(dummy_dataset, batch_size=10)
    calibrate_model(prepared_model, dummy_loader)
    print("Calibration complete.")

    # 5. 較正済みモデルをINT8に変換
    quantized_model = convert(prepared_model, inplace=False)

    # 6. 量子化済みモデルを保存
    torch.jit.save(torch.jit.script(quantized_model), "quantized_resnet50.pt")
    print("Quantized model saved as quantized_resnet50.pt")

if __name__ == '__main__':
    main()

# quantized_server.py
from fastapi import FastAPI
import torch
import time

app = FastAPI()

# 1. JITスクリプト化された量子化モデルを読み込む
model = torch.jit.load("quantized_resnet50.pt")
# JITモデルではmodel.eval()やtorch.inference_mode()は不要

@app.post("/predict")
async def predict():
    # CPU量子化モデルの場合、入力テンソルはCPUで作成する必要がある
    dummy_input = torch.randn(1, 3, 224, 224)

    start_time = time.perf_counter()
    output = model(dummy_input)
    end_time = time.perf_counter()

    latency_ms = (end_time - start_time) * 1000
    
    probabilities = torch.nn.functional.softmax(output, dim=0)
    top5_prob, top5_catid = torch.topk(probabilities, 5)

    return {
        "latency_ms": f"{latency_ms:.2f}",
        "top_category": int(top5_catid),
        "top_probability": float(top5_prob)
    }

{"latency_ms":"45.12","top_category":714,"top_probability":0.0071...}

# export_to_onnx.py
import torch
import onnxruntime

# 先に保存した量子化モデルを読み込む
quantized_model = torch.jit.load("quantized_resnet50.pt")
quantized_model.eval()

# モデルグラフをトレースするためのダミー入力を提供
dummy_input = torch.randn(1, 3, 224, 224)
onnx_path = "quantized_resnet50.onnx"

# モデルを表現
print("Exporting model to ONNX format...")
torch.onnx.export(
    quantized_model,
    dummy_input,
    onnx_path,
    export_params=True,
    opset_version=13, # 一般的で安定したバージョン
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
print(f"Model saved to {onnx_path}")

# オプション：ONNXモデルを検証
print("Verifying the ONNX model...")
ort_session = onnxruntime.InferenceSession(onnx_path)
ort_inputs = {ort_session.get_inputs().name: dummy_input.numpy()}
ort_outs = ort_session.run(None, ort_inputs)
print("ONNX model verified successfully.")

# onnx_server.py
from fastapi import FastAPI
import onnxruntime as ort
import numpy as np
import time

app = FastAPI()

# 1. ONNX Runtime推論セッションを作成
ort_session = ort.InferenceSession("quantized_resnet50.onnx")
input_name = ort_session.get_inputs().name

@app.post("/predict")
async def predict():
    # 2. ダミー入力を作成、今度はNumPy配列として
    dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

    # 3. ONNX Runtimeで推論を実行
    start_time = time.perf_counter()
    ort_inputs = {input_name: dummy_input}
    ort_outs = ort_session.run(None, ort_inputs)
    end_time = time.perf_counter()

    latency_ms = (end_time - start_time) * 1000
    
    # 4. 出力を処理（ONNXはNumPy配列を返す）
    output = ort_outs
    # softmaxとtop-kロジックは同じ、numpyを使用
    exp_output = np.exp(output)
    probabilities = exp_output / np.sum(exp_output)
    top5_catid = np.argsort(probabilities)[-5:][::-1]
    
    return {
        "latency_ms": f"{latency_ms:.2f}",
        "top_category": int(top5_catid),
    }

{"latency_ms":"28.31","top_category":714}

# batched_onnx_server.py
import asyncio
import time
from concurrent.futures import ThreadPoolExecutor
from queue import Queue, Empty
import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, HTTPException

# --- 設定 ---
MAX_BATCH_SIZE = 16
BATCH_TIMEOUT_MS = 5 # リクエストを収集する5msウィンドウ

# --- グローバルリソース ---
app = FastAPI()
request_queue = Queue()
# 結果を正しいリクエストに戻すための辞書
results_dict = {} 

# --- ONNXモデルロード ---
ort_session = ort.InferenceSession("quantized_resnet50.onnx")
input_name = ort_session.get_inputs().name

# --- バッチング用バックグラウンドワーカー ---
def inference_worker():
    while True:
        batch_items = []
        # 最初のリクエストを待つ
        request_id, input_data = request_queue.get()
        batch_items.append((request_id, input_data))

        # 時間ウィンドウ内でより多くのリクエストを収集
        start_time = time.perf_counter()
        while (time.perf_counter() - start_time) * 1000 < BATCH_TIMEOUT_MS and \
              len(batch_items) < MAX_BATCH_SIZE:
            try:
                # ノンブロッキングget
                next_id, next_data = request_queue.get_nowait()
                batch_items.append((next_id, next_data))
            except Empty:
                break # キューにこれ以上のアイテムがない

        # 推論用にバッチを準備
        batch_size = len(batch_items)
        if batch_size > 0:
            batched_input = np.concatenate([item for item in batch_items], axis=0)
            
            # バッチ推論を実行
            ort_inputs = {input_name: batched_input}
            ort_outs = ort_session.run(None, ort_inputs)
            
            # 結果を個々のリクエストに戻す
            for i in range(batch_size):
                request_id = batch_items[i]
                result = ort_outs[i]
                results_dict[request_id] = result

# --- FastAPIエンドポイント ---
@app.on_event("startup")
async def startup_event():
    # バックグラウンドワーカーを別スレッドで起動
    executor = ThreadPoolExecutor(max_workers=1)
    executor.submit(inference_worker)

@app.post("/predict")
async def predict():
    request_id = time.time() # このリクエストの一意なID
    input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

    # リクエストをキューに入れる
    request_queue.put((request_id, input_data))

    # 推論ワーカーからの結果を待つ
    while request_id not in results_dict:
        await asyncio.sleep(0.001) # 結果をポーリング

    # 結果を取得して削除
    result = results_dict.pop(request_id)
    
    # 前と同様に結果を処理して返す
    exp_output = np.exp(result)
    probabilities = exp_output / np.sum(exp_output)
    top_category = int(np.argmax(probabilities))
    
    return {"top_category": top_category}