PyTorch TorchScript / ONNX エクスポート

モデルのトレーニング完了後、モデルを本番環境にデプロイする必要があります。PyTorch は複数のモデルエクスポート方法を提供しており、そのうちTorchScriptとONNXが最も一般的に使用される2つの形式です。

このセクションでは、これら2つのエクスポート方法の原理、使用方法、ベストプラクティスを詳しく説明します。

モデルのエクスポートとは、PyTorch モデルを異なるプラットフォームや異なるフレームワークで実行できる形式に変換するプロセスです。これは、モデルのデプロイ、モバイル推論、フレームワーク間の移行などのシナリオにおいて非常に重要です。

1. TorchScript 基礎

1.1 TorchScript とは

TorchScript は PyTorch のシリアライズ形式であり、Python コードをスタンドアロンで実行可能な C++ 仮想マシンコードに変換できます。TorchScript プログラムは Python インタープリタのない環境でも実行できます。

TorchScript の主な特徴：

動的グラフを静的グラフに変換する
Python のサブセット構文をサポートする
C++ 環境で実行できる
モデルの構造とパラメータを保持する

1.2 2つの変換方法

TorchScript には、モデルを TorchScript に変換する2つの方法があります：

TorchScript トレース（Tracing）：モデルを実行して操作を記録し、静的計算グラフを生成する
TorchScript スクリプト（Scripting）：Python コードを直接分析し、TorchScript にコンパイルする

2. TorchScript トレース (Tracing)

2.1 基本的なトレース方法

例

import torch
import torch.nn as nn

# ── モデルを定義 ──────────────────────────────────────
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2)
self.fc = nn.Linear(128 * 8 * 8, 10)

def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = x.view(x.size(0), -1)
x = self.fc(x)
return x

# モデルインスタンスを作成
model = SimpleNet()
model.eval()

# サンプル入力
example_input = torch.randn(1, 3, 64, 64)

# ── 方法1：torch.jit.trace によるトレース ───────────────────
# モデルを実行して操作を記録し、TorchScript を作成
traced_model = torch.jit.trace(model, example_input)

print("トレース後のモデル:")
print(traced_model)

# モデルを保存
traced_model.save("simple_net_traced.pt")

# モデルを読み込み
loaded_model = torch.jit.load("simple_net_traced.pt")

# 読み込んだモデルで推論を実行
output = loaded_model(example_input)
print(f"出力形状: {output.shape}")

2.2 トレースの制限

トレース方法にはいくつかの制限があります：

実際に実行された操作のみを記録する
制御フロー（if、for など）は固定される
動的なサイズの入力では問題が発生する可能性がある

例

# トレースの制限の例

class DynamicModel(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
# 制御フローはトレース時に固定される
if x.sum() > 0:
return x * 2
else:
return x / 2

model = DynamicModel()
model.eval()

# トレース時、if 分岐はトレース時のパスに固定される
example_input = torch.tensor([1.0])
traced = torch.jit.trace(model, example_input)

# 入力が異なっていても、同じ分岐が実行される
print(traced(torch.tensor([1.0]))) # 2
print(traced(torch.tensor([-1.0]))) # 依然として 2 であり、-0.5 ではない

制御フローを含むモデルには、トレースではなく TorchScript スクリプト（Scripting）を使用する必要があります。

3. TorchScript スクリプト (Scripting)

3.1 基本的な使用方法

例

import torch

# ── @torch.jit.script デコレータを使用 ───────────────
@torch.jit.script
def scripted_function(x: torch.Tensor) -> torch.Tensor:
"""スクリプト方式で関数を変換"""
if x.sum() > 0:
return x * 2
else:
return x / 2

# スクリプト化された関数をテスト
input1 = torch.tensor([1.0, 2.0])
input2 = torch.tensor([-1.0, -2.0])

print(scripted_function(input1)) # [2., 4.]
print(scripted_function(input2)) # [-0.5, -1.]

# ── モデルのスクリプト化 ───────────────────────────────────
class ScriptableModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 10)

@torch.jit.export
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.relu(self.fc(x))

@torch.jit.export
def predict(self, x: torch.Tensor) -> torch.Tensor:
"""追加のエクスポート方法"""
out = self.forward(x)
return torch.argmax(out, dim=1)

model = ScriptableModel()

# モデルをスクリプト化
scripted_model = torch.jit.script(model)
print(scripted_model)

# 保存
scripted_model.save("scripted_model.pt")

3.2 複雑なモデルのスクリプト化

例

# より複雑なスクリプト化の例：条件付きモデル

class ConditionModel(nn.Module):
def __init__(self, num_classes: int):
super().__init__()
self.num_classes = num_classes
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
)
self.classifier = nn.Linear(64, num_classes)

def forward(self, x: torch.Tensor, use_softmax: bool = False) -> torch.Tensor:
"""
動的条件をサポートするモデル
"""
features = self.features(x)
logits = self.classifier(features)

if use_softmax:
return torch.softmax(logits, dim=1)
else:
return logits

def get_prediction(self, x: torch.Tensor) -> torch.Tensor:
"""ヘルパーメソッド"""
logits = self.forward(x, use_softmax=False)
return torch.argmax(logits, dim=1)

# スクリプト化
model = ConditionModel(num_classes=10)
scripted_model = torch.jit.script(model, example_inputs=(torch.randn(1, 3, 32, 32),))

# テスト
test_input = torch.randn(2, 3, 32, 32)
output1 = scripted_model(test_input, use_softmax=False)
output2 = scripted_model(test_input, use_softmax=True)

print(f"Logits 出力形状: {output1.shape}")
print(f"Softmax 出力形状: {output2.shape}")

4. ONNX エクスポート

4.1 ONNX 基礎

ONNX（Open Neural Network eXchange）は、オープンなニューラルネットワーク交換フォーマットであり、異なる深層学習フレームワーク間でのモデル変換をサポートしています。

ONNX の利点：

フレームワーク横断：PyTorch、TensorFlow、Caffe2 などでサポート
プラットフォーム横断：CPU、GPU、モバイル端末などの複数のプラットフォームをサポート
ハードウェア最適化：ONNX Runtime を利用して効率的な推論が可能
豊富なツール：多数の最適化ツールとデプロイソリューションがある

4.2 基本的な ONNX エクスポート

例

import torch
import torch.nn as nn
import torchvision

# ── モデルを定義 ──────────────────────────────────────
class ImageClassifier(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten()
)
self.classifier = nn.Linear(64, 10)

def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x

model = ImageClassifier()
model.eval()

# サンプル入力
example_input = torch.randn(1, 3, 32, 32)

# ── ONNX にエクスポート ──────────────────────────────────
output_path = "image_classifier.onnx"

torch.onnx.export(
model,
example_input,
output_path,
export_params=True, # モデルのパラメータをエクスポート
opset_version=14, # ONNX バージョン
do_constant_folding=True, # 定数折りたたみ最適化
input_names=['input'], # 入力テンソル名
output_names=['output'], # 出力テンソル名
dynamic_axes={
'input': {0: 'batch_size'}, # 動的バッチ次元
'output': {0: 'batch_size'}
}
)

print(f"モデルは {output_path} にエクスポートされました")

# エクスポートしたモデルを検証
import onnx
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("ONNX モデルの検証に成功しました！")

4.3 複雑なモデルのエクスポート

例

# 完全な ResNet モデルをエクスポート

import torchvision.models as models

# 事前学習済みモデルを読み込み
model = models.resnet18(pretrained=True)
model.eval()

# サンプル入力（ResNet 標準の 224x224）
example_input = torch.randn(1, 3, 224, 224)

# エクスポート
torch.onnx.export(
model,
example_input,
"resnet18.onnx",
export_params=True,
opset_version=14,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size', 2: 'height', 3: 'width'},
'output': {0: 'batch_size'}
}
)

print("ResNet18 をエクスポートしました")

# エクスポート時に特別な層を処理
# 特別な処理が必要な層には workaround を使用
class ModelWithSpecialOps(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 64, 3, padding=1)

def forward(self, x):
# nn.functional のエイリアスではなく F.interpolate を使用
x = self.conv(x)
x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
return x

model = ModelWithSpecialOps()
model.eval()
torch.onnx.export(
model,
torch.randn(1, 3, 32, 32),
"special_ops.onnx",
input_names=['input'],
output_names=['output'],
opset_version=14
)

5. エクスポートの検証と最適化

5.1 エクスポートしたモデルの検証

例

import numpy as np
import onnx
import onnxruntime as ort

def verify_onnx_model(onnx_path, pytorch_model, test_input):
"""
ONNX モデルと PyTorch モデルの出力の一致性を検証する
"""
# 1. ONNX モデルを確認
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("✓ ONNX モデル構造の検証に成功")

# 2. PyTorch 推論
pytorch_model.eval()
with torch.no_grad():
pytorch_output = pytorch_model(test_input)

# 3. ONNX Runtime 推論
ort_session = ort.InferenceSession(onnx_path)
ort_output = ort_session.run(None, {'input': test_input.numpy()})[0]

# 4. 出力を比較
diff = np.abs(pytorch_output.numpy() - ort_output).max()
print(f"✓ 最大出力差異: {diff:.6f}")

if diff < 1e-5:
print("✓ 出力検証に合格")
return True
else:
print("⚠ 出力差異が大きい")
return False

# 使用例
model = ImageClassifier()
model.eval()

# まずエクスポート
torch.onnx.export(
model,
torch.randn(1, 3, 32, 32),
"test_model.onnx",
input_names=['input'],
output_names=['output']
)

# 検証
test_input = torch.randn(2, 3, 32, 32)
verify_onnx_model("test_model.onnx", model, test_input)

5.2 ONNX Runtime の最適化

例

import onnxruntime as ort
from onnxruntime import GraphOptimizationLevel

# 推論セッションを作成し、最適化を構成する
def create_optimized_session(onnx_path, providers=None):
if providers is None:
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

# セッションオプションを構成
sess_options = ort.SessionOptions()

# グラフ最適化を有効化
sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL

# その他の最適化を有効化
sess_options.intra_op_num_threads = 4
sess_options.inter_op_num_threads = 4

# セッションを作成
session = ort.InferenceSession(onnx_path, sess_options, providers=providers)

return session

# 最適化されたセッションで推論を実行
session = create_optimized_session("resnet18.onnx")

# 入出力名を取得
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 推論
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = session.run([output_name], {input_name: input_data})[0]

print(f"推論出力形状: {output.shape}")

5.3 モデルの量子化

例

# ONNX モデルの量子化

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

def quantize_onnx_model(input_path, output_path):
"""
動的量子化 ONNX モデル
モデルサイズを削減し、推論を高速化する
"""
# 動的量子化（キャリブレーションデータ不要）
quantize_dynamic(
input_path,
output_path,
# 量子化する重みの型を指定
weight_type=QuantType.QInt8
)

# ファイルサイズの比較を取得
import os
original_size = os.path.getsize(input_path) / 1024 / 1024
quantized_size = os.path.getsize(output_path) / 1024 / 1024

print(f"元のモデル: {original_size:.2f} MB")
print(f"量子化モデル: {quantized_size:.2f} MB")
print(f"圧縮率: {original_size/quantized_size:.2f}x")

# 使用
quantize_onnx_model("resnet18.onnx", "resnet18_quantized.onnx")

6. モバイル端末へのデプロイ

6.1 モバイル形式へのエクスポート

例

# 方法1: TorchScript モバイル端末デプロイ
# 1. モデルをトレース
model = ImageClassifier()
model.eval()
example_input = torch.randn(1, 3, 32, 32)

traced_model = torch.jit.trace(model, example_input)

# 2. モデルを最適化（モバイル端末向け最適化）
optimized_model = torch.jit.optimize_for_inference(traced_model)

# 3. モバイル端末モデルを保存
optimized_model.save("mobile_model.pt")

# 方法2: torchmobile を使う（利用可能な場合）
# torchmobile はより軽量なモバイル端末デプロイに使用
# クロスコンパイルは公式ドキュメントを参照してください

# 方法3: ONNX にエクスポートしてモバイルランタイムを使用
# iOS: Core ML Tools を使用
# Android: ONNX Runtime Mobile を使用

print("モバイル端末モデルの準備ができました")

6.2 Core ML エクスポート（iOS）

iOS プラットフォームでは、Core ML Tools を使用してモデルを Core ML 形式に変換できます。まず ONNX 形式に変換してから、Core ML Tools で変換する必要があります。

例

# coremltools のインストールが必要
# pip install coremltools

# Core ML エクスポート手順
# 注意: これは macOS でのみ利用可能

# 手順1: PyTorch モデルを ONNX にエクスポート
# from pytorch_example import your_model
# model = your_model()
# model.eval()
# example_input = torch.randn(1, 3, 224, 224)
# torch.onnx.export(
# model,
# example_input,
# "model.onnx",
# input_names=['input'],
# output_names=['output']
# )

# 手順2: Core ML Tools を使用して Core ML 形式に変換
# from coremltools.converters import onnx as onnx_coreml
# coreml_model = onnx_coreml.convert(
# model="model.onnx",
# minimum_deployment_target='13',
# image_input_names=['input']
# )
# coreml_model.save("model.mlmodel")

# 使用例
print("iOS デプロイには macOS 環境が必要です")
print("詳細手順：1. pip install coremltools 2. ONNX へのエクスポート 3. coremltools を使用した変換")

7. ベストプラクティスとよくある質問

7.1 エクスポートの問題トラブルシューティング

問題	原因	解決策
エクスポート失敗	サポートされていない操作	より高いバージョンの opset_version を使用するか、操作を置き換える
出力の不一致	動的制御フロー	トレースではなくスクリプト化を使用するか、入力サイズを固定する
ONNX Runtime エラー	オペレーターがサポートされていない	ONNX がサポートする操作リストを確認

7.2 エクスポート形式の比較

形式	長所	短所	適用シナリオ
TorchScript (.pt)	PyTorch ネイティブサポート	クロスフレームワークのサポートが弱い	デスクトップ/サーバーデプロイ
ONNX (.onnx)	クロスフレームワーク、クロスプラットフォーム	一部の操作がサポートされていない	汎用デプロイ

7.3 エクスポートチェックリスト

例

# エクスポート前チェックリスト

def export_checklist(model, example_input):
"""エクスポート前のチェック"""
model.eval()

# 1. モデルが推論モードであることを確認
print(f"モデルモード: {'eval' if not model.training else 'train'}")

# 2. 入力サイズを確認
print(f"サンプル入力形状: {example_input.shape}")

# 3. 順伝播を検証
with torch.no_grad():
output = model(example_input)
print(f"出力形状: {output.shape}")

# 4. シリアライズ不可能な操作がないかを確認
# 例: lambda 関数
for name, module in model.named_modules():
if hasattr(module, '__call__'):
pass # カスタムモジュールを確認

print("✓ エクスポートチェック完了")

# 使用例
model = ImageClassifier()
example_input = torch.randn(1, 3, 32, 32)
export_checklist(model, example_input)

7.4 完全なエクスポートフロー

例

# モデルエクスポートの完全なフロー例

class ProductionModel(nn.Module):
"""本番向けモデル"""
def __init__(self):
super().__init__()
self.backbone = nn.Sequential(
nn.Conv2d(3, 64, 7, stride=2, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(3, stride=2, padding=1),
# ... その他のレイヤー
)
self.head = nn.Linear(512, 10)

def forward(self, x, return_features=False):
features = self.backbone(x)
features = torch.nn.functional.adaptive_avg_pool2d(features, 1).flatten(1)

if return_features:
return features

return self.head(features)

# 完全なエクスポートフロー
model = ProductionModel()
model.load_state_dict(torch.load("weights.pth")) # 重みを読み込み
model.eval()

# サンプル入力を準備
example_input = torch.randn(1, 3, 224, 224)

# 1. TorchScript をトレース
traced = torch.jit.trace(model, example_input)
traced.save("model_traced.pt")

# 2. ONNX にエクスポート
torch.onnx.export(
traced,
example_input,
"model.onnx",
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
opset_version=14
)

# 3. ONNX を最適化
# onnx-simplifier を使用して冗長な操作を除去
# onnxsim model.onnx model_simplified.onnx

print("すべての形式のエクスポートが完了しました！")

その他の拡張