AI Edge Torch – Inferência de alto desempenho de modelos do PyTorch em dispositivos móveis

MAI 14, 2024

Cormac Brick Principal Engineer

Advait Jain Software Engineer

Haoliang Zhang Software Engineer

Temos o prazer de anunciar o Google AI Edge Torch, um caminho direto do PyTorch para o ambiente de execução do TensorFlow Lite (TFLite), com excelente abrangência de modelos e desempenho da CPU. O TFLite já trabalha com modelos escritos em Jax, Keras e TensorFlow, e agora estamos adicionando o PyTorch como parte de um compromisso mais amplo de fornecer várias opções de framework.

Agora, essa nova oferta está disponível como parte do Google AI Edge, um conjunto de ferramentas com acesso fácil a tarefas de ML prontas para uso, frameworks que permitem criar pipelines de ML e executar LLMs populares e modelos personalizados – tudo no dispositivo. Esta é a primeira de uma série de postagens do blog sobre os lançamentos do Google AI Edge que ajudarão os desenvolvedores a criar recursos com tecnologia de IA e implantá-los em várias plataformas com facilidade.

O AI Edge Torch foi lançado hoje em versão Beta com:

Integração direta com o PyTorch
Excelente desempenho de CPU e suporte inicial a GPU
Validado em mais de 70 modelos de torchvision, timm, torchaudio e HuggingFace
Suporte a mais de 70% dos operadores core_aten no PyTorch
Compatibilidade com o ambiente de execução atual do TFLite, sem necessidade de alteração no código de implantação
Suporte à visualização do Model Explorer em vários estágios do fluxo de trabalho.

Uma experiência simples e centrada no PyTorch

O Google AI Edge Torch foi criado do zero para fornecer uma ótima experiência à comunidade do PyTorch, com APIs que parecem nativas e oferecem um caminho de conversão fácil.

import torchvision
import ai_edge_torch
 
# Initialize model
resnet18 = torchvision.models.resnet18().eval()
 
# Convert
sample_input = (torch.randn(4, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18, sample_input)
 
# Inference in Python
output = edge_model(*sample_input)
 
# Export to a TfLite model for on-device deployment
edge_model.export('resnet.tflite'))

Nos bastidores, ai_edge_torch.convert() é integrado ao TorchDynamo usando torch.export, que é a maneira do PyTorch 2.x de exportar modelos do PyTorch para representações de modelos padronizadas destinadas a serem executadas em diferentes ambientes. Nossa implementação atual tem suporte a mais de 60% dos operadores core_aten, e planejamos aumentar esse número consideravelmente à medida que avançarmos para a versão 1.0 de ai_edge_torch. Incluímos exemplos que mostram a quantização PT2E, a estratégia de quantização nativa do PyTorch2, para permitir fluxos de trabalho de quantização facilitados. Mal podemos esperar para ouvir a comunidade do PyTorch a fim de encontrar formas de melhorar a experiência dos desenvolvedores ao levar a inovação que começa no PyTorch para um amplo conjunto de dispositivos.

Abrangência e desempenho

Antes desta versão, muitos desenvolvedores estavam usando caminhos fornecidos pela comunidade, como ONNX2TF, para possibilitar o uso de modelos do PyTorch no TFLite. Nosso objetivo no desenvolvimento do AI Edge Torch era reduzir a dificuldade para o desenvolvedor, fornecer uma ótima abrangência de modelos e continuar nossa missão de oferecer desempenho de primeiro nível em dispositivos Android.

Quanto à abrangência, nossos testes demonstram melhorias significativas no conjunto definido de modelos para os fluxos de trabalho existentes, particularmente ONNX2TF

Table showing performance improvement in existing workflows over defined set models

Quanto ao desempenho, nossos testes mostram um desempenho consistente com a linha de base do ONNX2TF, e também mostram um desempenho consideravelmente melhor do que o ambiente de execução do ONNX:

Table showing performance with ONNX2TF baseline

Veja o desempenho detalhado por modelo para o subconjunto dos modelos abrangidos pelo ONNX:

Chart showing per model TFLite latency relative to ONNX

Figura – Latência de inferência por rede em comparação ao ONNX, medida no Pixel8, precisão fp32, XNNPACK fixado em 4 threads para auxiliar a reprodutibilidade, média de 100 execuções após uma fase de preparação com 20 iterações

Adoção antecipada e parcerias

Nos últimos meses, trabalhamos em estreita colaboração com parceiros de adoção antecipada, incluindo Shopify, Adobe e Niantic, para melhorar nosso suporte ao PyTorch. O ai_edge_torch já está sendo usado pela equipe da Shopify para realizar a remoção do plano de fundo em imagens de produtos, feita no dispositivo, e estará disponível em uma próxima versão do aplicativo da Shopify.

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

Parcerias para componentes eletrônicos e delegates

Também trabalhamos em estreita colaboração com parceiros para trabalhar no suporte de hardware em CPUs, GPUs e aceleradores, incluindo ARM, Google Tensor G3, MediaTek, Qualcomm e Samsung System LSI. Por meio dessas parcerias, melhoramos o desempenho e a abrangência, e também validamos os arquivos do TFLite gerados pelo PyTorch nos delegates do acelerador.

Também é um grande prazer anunciar o novo delegate da Qualcomm para o TensorFlow Lite, que já está disponível abertamente aqui para qualquer desenvolvedor usar. Os delegates do TFLite são módulos de software complementares que ajudam a acelerar a execução em GPUs e em aceleradores de hardware. Este novo delegate QNN é compatível com a maioria dos modelos em nosso conjunto de testes Beta do PyTorch. Além disso, ele também tem suporte a um amplo conjunto de componentes eletrônicos da Qualcomm e oferece acelerações médias consideráveis em relação à CPU (20 vezes) e GPU (5 vezes) utilizando as unidades de processamento neural e o DSP da Qualcomm. Para facilitar o teste, a Qualcomm também lançou recentemente seu novo AI Hub. O Qualcomm AI Hub é um serviço em nuvem que permite aos desenvolvedores testar modelos do TFLite em um amplo conjunto de dispositivos Android, além de fornecer visibilidade dos ganhos de desempenho disponíveis em diferentes dispositivos usando o delegate QNN.

O que vem a seguir?

Nos próximos meses, continuaremos fazendo iterações abertamente, com lançamentos que vão expandir a abrangência de modelos, melhorar o suporte à GPU e possibilitar novos modos de quantização à medida que prosseguirmos para a versão 1.0. Na parte 2 desta série, daremos uma olhada mais aprofundada na API AI Edge Torch Generative, que permite aos desenvolvedores levar modelos de GenAI personalizados para a borda, com um ótimo desempenho.

Gostaríamos de agradecer a todos os nossos clientes de acesso antecipado por seus valiosos comentários que nos ajudaram a detectar bugs precoces e a garantir uma experiência tranquila para os desenvolvedores. Também gostaríamos de agradecer aos parceiros de hardware e aos colaboradores do ecossistema do XNNPACK, que nos ajudaram a melhorar o desempenho em uma ampla gama de dispositivos. Também gostaríamos de agradecer à comunidade do PyTorch em geral pela orientação e apoio.

Agradecimentos

^{Gostaríamos de agradecer a cada membro da equipe que contribuiu para este trabalho: 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 e a equipe StableHLO.}