pytorch-deployment

PyTorch - デプロイと本番環境エンジニアリング

高性能な環境にモデルをデプロイするには、多くの場合、Python への依存を取り除く必要があります。このガイドでは、C++ で読み込める形式へのモデルのシリアライズ、エッジデバイス向けの最適化、TensorRT のような高スループット推論エンジンでの実行方法を説明します。

使用する場面

• Jupyter Notebook のモデルを本番環境の Web サーバー（FastAPI/Go/Rust）に移行する場合。
• ニューラルネットワークを C++ アプリケーション（LibTorch）に組み込む場合。
• モバイルデバイス（iOS/Android）やエッジハードウェア（NVIDIA Jetson）で推論を実行する場合。
• 専用ハードウェアバックエンド（OpenVINO、TensorRT）を使って推論速度を向上させる場合。
• 異なるバージョンの PyTorch 間でモデルの再現性を確保する場合。

基本原則

1. スクリプト化とトレース

• トレース：PyTorch は「サンプルデータ」を使ってモデルを一度実行し、すべての演算を記録します。高速ですが、Python の制御フロー（if、for）は無視します。
• スクリプト化：PyTorch はモジュールの Python ソースコードをコンパイルします。準備には時間がかかりますが、ロジックと制御フローが保持されます。

2. ONNX ブリッジ

ONNX（Open Neural Network Exchange）はクロスプラットフォーム形式です。ONNX にエクスポートしたモデルは、推論時に標準の PyTorch より高速なことが多い、Microsoft の ONNX Runtime で実行できます。

3. 量子化

重みを float32（4 バイト）から int8（1 バイト）に削減します。これによりモデルサイズは 4 分の 1 になり、CPU では推論が 2～3 倍高速化することがあります。

クイックリファレンス：エクスポートパターン

import torch

model = MyModel().eval()
example_input = torch.randn(1, 3, 224, 224)

# 1. Tracing (Most common)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_jit.pt")

# 2. Scripting (For dynamic logic)
scripted_model = torch.jit.script(model)
scripted_model.save("model_script.pt")

# 3. ONNX Export
torch.onnx.export(model, example_input, "model.onnx", 
                  input_names=['input'], output_names=['output'],
                  dynamic_axes={'input': {0: 'batch_size'}})

重要なルール

✅ 推奨事項

• エクスポート前に model.eval() を呼び出す - これにより BatchNorm レイヤーと Dropout レイヤーが固定されます。呼び出し忘れると、予測が不正確になります。
• torch.no_grad() を使う - 不要な勾配追跡メタデータが保存されないよう、エクスポート処理は必ず no_grad コンテキストで囲みます。
• ONNX で dynamic_axes を定義する - モデルが異なるバッチサイズや画像解像度を扱う場合、エクスポート時に必ず指定します。
• エクスポートの精度を検証する - torch.allclose() を使って、元の Python モデルとエクスポートしたモデルの出力を必ず比較します。
• Python でのデプロイには torch.compile を使う - Python 内でデプロイする場合、パフォーマンス向上のため JIT の代わりに torch.compile（PyTorch 2.0+）を使います。

❌ 避けるべき事項

• if/else を含むモデルに JIT トレースを使わない - トレーサーはサンプル実行時に選択された分岐だけを記録します。
• 通常、前処理をモデルに含めない - TorchScript 互換の演算を使う場合を除き、柔軟性を高めるため、画像のリサイズや正規化はコアモデルの外で行います。
• 量子化の警告を無視しない - カスタム活性化関数など、一部のレイヤーは int8 をサポートせず、float32 にフォールバックするため、効果が低下します。

高度な最適化

学習後の量子化（静的）

import torch.quantization

# 1. Set backend (x86 or ARM)
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 2. Prepare and Calibrate (Run some data through the model)
model_prepared = torch.quantization.prepare(model)
# ... run calibration loop ...

# 3. Convert
model_int8 = torch.quantization.convert(model_prepared)

LibTorch（C++ デプロイ）

C++ で TorchScript モデルを読み込むには：

#include <torch/script.h>

int main() {
    // Load model
    torch::jit::script::Module module = torch::jit::load("model_jit.pt");
    
    // Create input tensor
    auto input = torch::randn({1, 3, 224, 224});
    
    // Run inference
    at::Tensor output = module.forward({input}).toTensor();
    std::cout << output.slice(1, 0, 5) << std::endl;
}

実践的なワークフロー

1. モバイル向けの最適化（Lite Interpreter）

モバイルにデプロイする場合、標準の TorchScript は重すぎます。「Mobile」オプティマイザーを使います。

from torch.utils.mobile_optimizer import optimize_for_mobile
optimized_model = optimize_for_mobile(traced_model)
optimized_model._save_for_lite_interpreter("model_mobile.ptl")

2. ONNX Runtime によるデプロイ

import onnxruntime as ort

session = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])
outputs = session.run(None, {"input": example_input.numpy()})

よくある落とし穴と解決策

JIT での「Missing Attribute」エラー

TorchScript からは、初期化後にモデルへ追加された属性が認識できません。

# ✅ Solution: Define all needed attributes in __init__ or use @torch.jit.export

動的シェイプの失敗

モデルが計算で x.shape[0] を使用している場合、トレースによってその値がハードコードされることがあります。

# ✅ Solution: Use Scripting or ensure calculations use tensor methods 
# like .size(0) which JIT understands.

PyTorch のデプロイは、科学と現実世界をつなぐ架け橋です。これらのツールを習得すれば、発見がノートブックにとどまらず、次世代のインテリジェントシステムを動かす力になります。