PyTorch モデルを TF Lite に変換する

AI Edge Torch は、PyTorch モデルを .tflite 形式に変換できるライブラリで、TensorFlow Lite と MediaPipe でこれらのモデルを実行できます。これは、モデルを完全にデバイス上で実行するモバイルアプリを作成するデベロッパーにとって特に便利です。AI Edge Torch は幅広い CPU をカバーしており、初期の GPU と NPU をサポートしています。

PyTorch モデルを TF Lite に変換するには、Pytorch コンバータのクイックスタートを使用します。詳細については、AI Edge Torch の GitHub リポジトリをご覧ください。

大規模言語モデル（LLM）や Transformer ベースのモデルを特に変換する場合は、Generative Torch API を使用します。この API は、モデルの作成や量子化など、Transformer 固有の変換の詳細を処理します。

コンバージョンのワークフロー

次の手順では、PyTorch モデルを TensorFlow Lite にエンドツーエンドで変換する簡単な方法を示します。

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 モデルを変換したら、変換された新しい TF Lite モデルで推論を実行できます。

output = edge_model(*sample_inputs)

変換後のモデルは、後で使用するために .tflite 形式でエクスポートして保存できます。

edge_model.export('resnet.tflite')