Mengonversi model PyTorch ke TF Lite

AI Edge Torch adalah library yang memungkinkan Anda mengonversi model PyTorch ke dalam format .tflite, sehingga Anda dapat menjalankan model tersebut dengan TensorFlow Lite dan MediaPipe. Hal ini sangat membantu bagi developer yang membuat aplikasi seluler yang menjalankan model sepenuhnya di perangkat. AI Edge Torch menawarkan cakupan CPU yang luas, dengan dukungan GPU dan NPU awal.

Untuk mulai mengonversi model PyTorch ke TF Lite, gunakan panduan memulai pengonversi Pytorch. Untuk mengetahui informasi selengkapnya, lihat repositori GitHub AI Edge Torch.

Jika Anda secara khusus mengonversi Model Bahasa Besar (LLM) atau model berbasis transformer, gunakan Generative Torch API, yang menangani detail konversi khusus transformer seperti penulisan dan kuantisasi model.

Alur kerja konversi

Langkah-langkah berikut menunjukkan konversi end-to-end yang sederhana dari model PyTorch ke TensorFlow Lite.

Impor Obor Edge AI

Mulai dengan mengimpor paket pip AI Edge Torch (ai-edge-torch), bersama PyTorch.

import ai_edge_torch
import torch

Untuk contoh ini, kami juga memerlukan paket berikut:

import numpy
import torchvision

Menginisialisasi dan mengonversi model

Kita akan mengonversi ResNet18, model pengenalan citra yang populer.

resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()

Gunakan metode convert dari library AI Edge Torch untuk mengonversi model PyTorch.

sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)

Menggunakan model

Setelah mengonversi model Pytorch, Anda dapat menjalankan inferensi dengan model TF Lite yang dikonversi baru.

output = edge_model(*sample_inputs)

Anda dapat mengekspor dan menyimpan model yang dikonversi dalam format .tflite untuk digunakan di masa mendatang.

edge_model.export('resnet.tflite')