AI Edge Torch: Inferensi Model Pytorch Berkinerja Tinggi pada Perangkat Seluler

MEI 14, 2024

Cormac Brick Principal Engineer

Advait Jain Software Engineer

Haoliang Zhang Software Engineer

Kami sangat senang bisa mengumumkan Google AI Edge Torch - jalur langsung dari PyTorch ke runtime TensorFlow Lite (TFLite) dengan cakupan model dan performa CPU yang luar biasa. TFLite sudah dapat digunakan dengan model yang ditulis dalam Jax, Keras, dan TensorFlow, dan sekarang kami menambahkan PyTorch sebagai bagian dari komitmen besar kami terhadap opsionalitas framework.

Penawaran baru ini sekarang tersedia sebagai bagian dari Google AI Edge, rangkaian alat dengan akses mudah ke tugas ML siap pakai, framework yang memungkinkan Anda membangun pipeline ML, dan menjalankan LLM serta model khusus populer – semuanya di perangkat. Ini adalah seri pertama dari postingan blog yang membahas tentang rilis Google AI Edge yang akan membantu developer membangun fitur-fitur berteknologi AI, dan menerapkannya dengan mudah di berbagai platform.

AI Edge Torch dirilis dalam versi Beta hari ini dengan menghadirkan:

Integrasi PyTorch langsung
Performa CPU yang luar biasa dan dukungan GPU awal
Divalidasi pada lebih dari 70 model dari torchvision, timm, torchaudio, dan HuggingFace
Dukungan untuk > 70% operator core_aten di PyTorch
Kompatibilitas dengan runtime TFLite saat ini, tanpa perlu mengubah kode deployment
Dukungan untuk visualisasi Model Explorer pada beberapa tahap alur kerja.

Pengalaman yang berpusat pada PyTorch dan simpel

Google AI Edge Torch dibangun dari dasar untuk memberikan pengalaman yang luar biasa kepada komunitas PyTorch, dengan API yang terasa native, dan menyediakan jalur konversi yang mudah.

import torchvision
import ai_edge_torch
 
# Inisialisasi model
resnet18 = torchvision.models.resnet18().eval()
 
# Konversi
sample_input = (torch.randn(4, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18, sample_input)
 
# Inferensi di Python
output = edge_model(*sample_input)
 
# Ekspor ke model TfLite untuk deployment di perangkat
edge_model.export('resnet.tflite'))

Di belakang layar, ai_edge_torch.convert() terintegrasi dengan TorchDynamo menggunakan torch.export - yang merupakan cara PyTorch 2.x untuk mengekspor model PyTorch ke dalam representasi model terstandarisasi yang dirancang untuk berjalan di lingkungan yang berbeda. Implementasi kami saat ini mendukung lebih dari 60% operator core_aten, yang ingin kami tingkatkan secara signifikan sembari kami membangun rilis 1.0 ai_edge_torch. Kami telah menyertakan contoh yang menunjukkan kuantisasi PT2E, pendekatan kuantisasi native untuk PyTorch2, untuk memudahkan alur kerja kuantisasi. Kami sangat senang mendengar masukan dari komunitas PyTorch untuk menemukan cara meningkatkan pengalaman developer ketika membawa inovasi yang dimulai di PyTorch ke berbagai perangkat.

Cakupan & Performa

Sebelum rilis ini, banyak developer menggunakan jalur yang disediakan komunitas seperti ONNX2TF untuk mengaktifkan model PyTorch di TFLite. Tujuan kami mengembangkan AI Edge Torch adalah untuk mengurangi gesekan antar developer, memberikan cakupan model yang besar, dan melanjutkan misi kami untuk memberikan performa terbaik di kelasnya pada perangkat Android.

Dari sisi cakupan, pengujian kami menunjukkan peningkatan yang signifikan pada rangkaian model yang telah ditentukan dibandingkan alur kerja saat ini, khususnya ONNX2TF

Table showing performance improvement in existing workflows over defined set models

Dari sisi performa, pengujian kami menunjukkan performa yang konsisten dengan baseline ONNX2TF, sekaligus menunjukkan performa yang jauh lebih baik daripada runtime ONNX:

Table showing performance with ONNX2TF baseline

Ini menunjukkan performa tiap model secara terperinci pada subset model yang dicakup ONNX:

Chart showing per model TFLite latency relative to ONNX

Gambar: Latensi inferensi tiap jaringan dibandingkan dengan ONNX, diukur pada Pixel8, presisi fp32, XNNPACK disetel ke 4 thread untuk membantu reproduktifitas, rerata dari 100 kali berjalan setelah 20 iterasi pemanasan

Adopsi Awal dan Kemitraan

Dalam beberapa bulan terakhir, kami bekerja sama dengan mitra adopsi awal termasuk Shopify, Adobe, dan Niantic untuk meningkatkan dukungan PyTorch kami. ai_edge_torch telah digunakan oleh tim di Shopify untuk menghapus latar belakang gambar produk di perangkat dan akan tersedia dalam rilis aplikasi Shopify mendatang.

Quote image with text reads "Converting PyTorch models to run locally on Android was complex. Google's new tools simplify this, enabling fast creation of mobile-ready PyTorch models - Mustapha Ali, Shopify, Director of Engineering

Kemitraan silikon & delegasi

Kami juga bekerja sama dengan para mitra untuk mengerjakan dukungan hardware untuk semua CPU, GPU, dan akselerator - termasuk Arm, Google Tensor G3, MediaTek, Qualcomm, dan Samsung System LSI. Melalui kemitraan ini, kami meningkatkan performa serta cakupan, dan telah memvalidasi file TFLite yang dihasilkan PyTorch pada delegasi akselerator.

Kami juga sangat gembira bisa mengumumkan delegasi TensorFlow Lite baru dari Qualcomm, yang kini tersedia secara terbuka di sini untuk digunakan oleh setiap developer. Delegasi TFLite adalah modul software add-on yang membantu mempercepat eksekusi pada GPU dan akselerator hardware. Delegasi QNN baru ini mendukung sebagian besar model dalam rangkaian pengujian PyTorch Beta kami, sekaligus memberikan dukungan untuk berbagai macam silikon Qualcomm, serta memberikan peningkatan kecepatan rata-rata yang signifikan dibandingkan dengan CPU (20x) dan GPU (5x) dengan memanfaatkan DSP Qualcomm dan unit pemrosesan neural. Untuk memudahkan pengujian, Qualcomm baru saja merilis AI Hub terbaru mereka. Qualcomm AI Hub adalah layanan cloud yang memungkinkan developer menguji model TFLite pada banyak perangkat Android, dan menyediakan visibilitas peningkatan performa yang tersedia pada beragam perangkat menggunakan delegasi QNN.

Apa berikutnya?

Dalam beberapa bulan ke depan, kami akan terus melakukan iterasi secara terbuka, dengan rilis yang memperluas cakupan model, meningkatkan dukungan GPU, dan mengaktifkan mode kuantisasi baru sembari membangun rilis 1.0. Pada bagian 2 seri ini, kita akan melihat lebih mendalam tentang AI Edge Torch Generative API, yang memungkinkan developer menghadirkan model GenAI khusus ke edge dengan performa yang luar biasa.

Kami ingin mengucapkan terima kasih kepada semua pengguna akses awal kami atas masukan berharga mereka yang membantu kami menemukan bug dan memastikan pengalaman developer yang mulus. Kami juga ingin berterima kasih kepada mitra hardware, dan kontributor ekosistem untuk XNNPACK yang telah membantu kami meningkatkan performa pada beragam perangkat. Kami juga ingin berterima kasih kepada komunitas Pytorch yang lebih luas atas bimbingan dan dukungan mereka.

Ucapan Terima kasih

^{Kami ingin mengucapkan terima kasih kepada semua anggota tim yang telah berkontribusi dalam pekerjaan ini: Aaron Karp, Advait Jain, Akshat Sharma, Alan Kelly, Arian Arfaian, Chun-nien Chan, Chuo-Ling Chang, Claudio Basille, Cormac Brick, Dwarak Rajagopal, Eric Yang, Gunhyun Park, Han Qi, Haoliang Zhang, Jing Jin, Juhyun Lee, Jun Jiang, Kevin Gleason, Khanh LeViet, Kris Tonthat, Kristen Wright, Lu Wang, Luke Boyer, Majid Dadashi, Maria Lyubimtseva, Mark Sherwood, Matthew Soulanille, Matthias Grundmann, Meghna Johar, Milad Mohammadi, Na Li, Paul Ruiz, Pauline Sho, Ping Yu, Pulkit Bhuwalka, Ram Iyengar, Sachin Kotwani, Sandeep Dasgupta, Sharbani Roy, Shauheen Zahirazami, Siyuan Liu, Vamsi Manchala, Vitalii Dziuba, Weiyi Wang, Wonjoo Lee, Yishuang Pang, Zoe Wang, dan tim StableHLO.}

diposting di: