TorchScript によるモデルのエクスポート

PyTorch Lightning によるスケーラブルな AI モデルの構築

Sergiy Tkachuk

Director, GenAI Productivity

TorchScript とは

Python に依存しない
本番環境で高効率
例：
- モバイルデバイスへのデプロイ

スマートフォン

TorchScript とは

Python に依存しない
本番環境で高効率
例：
- モバイルデバイスへのデプロイ
- 本番システムでの高性能な推論

本番環境を表すギアを持つスマートフォンとノートパソコン

モデルを TorchScript に変換する

変換には2つの方法があります。

torch.jit.trace：入力例を使って実行をトレース
torch.jit.script：Python コードを解析してコンパイル

使い分け：

シンプルなモデルには trace を使用
制御フロー（ループなど）を含むモデルには script を使用

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted_model = torch.jit.script(model)

TorchScript モデルの保存と読み込み

モデルの保存：
- torch.jit.save：スクリプト化したモデルをファイルに保存
モデルの読み込み：
- torch.jit.load：推論のためにモデルを再読み込み

# Save the model
torch.jit.save(scripted_mod,"model.pt")

# Load the model
loaded_model=torch.jit.load("model.pt")

TorchScript で推論を実行する

手順：
- TorchScript モデルを読み込む
- 予測のためにモデルへ入力を渡す
- 出力は PyTorch の予測結果と同一

入力例：

入力テンソル：[1.0, 2.0, 3.0]

出力例：

出力テンソル：[2.0, 4.0, 6.0]

# Perform inference
input_arr = [1.0, 2.0, 3.0]
input_tensor = torch.tensor(input_arr)

output = loaded_model(input_tensor)
print(output)

TorchScript のまとめ

torch.jit.trace：静的なモデルに対応
torch.jit.script：動的な制御フローを処理
torch.jit.save：スクリプト化したモデルを保存
torch.jit.load：推論のために再読み込み

練習しましょう！

PyTorch Lightning によるスケーラブルな AI モデルの構築