General Compilation and Mixed-Precision Partitioning: A Combined Approach for Adaptive On-Device Learning

The application of machine learning (ML) is becoming more widespread, with a growing number of use cases. The development of centralized data training and the exponential growth of data generation raise significant privacy and security concerns. To mitigate these issues, on-device transfer learning (TL) offers the ability for models to adapt to local data without relying on a cloud connection. TL can leverage the knowledge gained from pre-trained models, hence, accelerating the training process. In parallel, ML compilers have become an essential tool for deploying ML models on a wide range of hardware platforms. However, current ML compilers focus on optimizing inference, neglecting on-device TL that can address real-world challenges, such as model performance degradation and distribution shift. This article proposes a mixed-precision partitioning algorithm, which identifies optimal partitioning layers for retraining, and a general compilation approach, based on Apache TVM, that enables on-device TL for embedded systems. We provide performance measurements for embedded graphics processing units (GPUs), central processing units (CPUs), and an artificial intelligence (AI) accelerator, demonstrating a simplified deployment process for deep neural networks (DNNs) with integrated on-device training capabilities.