AI Edge Torch は、PyTorch モデルを .tflite 形式に変換し、LiteRT と MediaPipe でモデルを実行できるようにするライブラリです。これは、モデルを完全にデバイス上で実行するモバイルアプリを作成するデベロッパーにとって特に便利です。AI Edge Torch は、幅広い CPU をサポートし、初期の GPU と NPU をサポートしています。
PyTorch モデルを LiteRT に変換するには、Pytorch コンバータのクイックスタートを使用します。詳細については、AI Edge Torch の GitHub リポジトリをご覧ください。
大規模言語モデル（LLM）またはトランスフォーマー ベースのモデルを変換する場合は、Generative Torch API を使用します。これは、モデルの作成や量子化など、トランスフォーマー固有の変換の詳細を処理します。
コンバージョン ワークフロー
次の手順では、PyTorch モデルから LiteRT への単純なエンドツーエンドの変換を示します。
AI Edge Torch をインポートする
まず、PyTorch とともに AI Edge Torch（ai-edge-torch）pip パッケージをインポートします。
import ai_edge_torch
import torch
この例では、次のパッケージも必要です。
import numpy
import torchvision
モデルを初期化して変換する
人気のある画像認識モデルである ResNet18 を変換します。
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()
AI Edge Torch ライブラリの convert メソッドを使用して、PyTorch モデルを変換します。
sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)
モデルを使用する
Pytorch モデルを変換したら、変換された新しい LiteRT モデルで推論を実行できます。
output = edge_model(*sample_inputs)
変換されたモデルは、.tflite 形式でエクスポートして保存し、後で使用できます。
edge_model.export('resnet.tflite')