PyTorch s’ouvre à l’API Android Neural Networks

PyTorch NNAPI

PyTorch Mobile introduit des fonctionnalités expérimentales centrées sur l’accélération matérielle. Dont une jonction avec l’API Neural Networks d’Android.

Connexion établie entre PyTorch et l’API Android Neural Networks (NNAPI). La jonction n’est encore que partielle. Elle s’inscrit dans une série de fonctionnalités expérimentales destinées à tirer parti de l’accélération matérielle pour l’inférence sur les terminaux mobiles.

Sur la partie GPU, les API Vulkan (Android) et Metal (iOS) sont mises à contribution. La passerelle avec la NNAPI vise les autres accélérateurs - essentiellement les NPU et les DSP. Facebook - créateur de PyTorch - a mené des tests avec un modèle intégré dans l'application Messenger. En l'occurrence, celui qui gère les arrière-plans immersifs (voir ci-dessous).

Sur un smartphone Google Pixel 5, le recours à la NNAPI a réduit de moitié la latence mesurée avec un cœur CPU. Sur un Pixel 4, il l'a divisée par près de quatre. Moyennant une sollicitation plus importante du processeur.

Pixel 3 NNAPI

Le différentiel est plus significatif sur l'architecture MobileNetV2 (vision par ordinateur).

MobileNetV2

L'utilisation de ce modèle avec la NNAPI a requis quelques étapes de préparation, listées dans ce tutoriel. De manière générale, quelques adaptations sont nécessaires. PyTorch présente effectivement des différences sémantiques avec la NNAPI. Par exemple sur l'organisation des poids synaptiques en mémoire. Ou sur la représentation des opérations de suréchantillonnage.

Pour ce « premier jet », le tandem PyTorch-NNAPI ne fonctionne qu'à partir d'Android 10. La prise en charge d'Android 8 et d'Android 9 est sur la feuille de route, sans échéance définie. Même chose pour les modèles qui exploitent le contrôle de flux ou encore le fallback automatique sur CPU.

Illustration principale © ProStokStudio - shutterstock.com