Tensor/IO for Android

Motivation and Introduction

Data privacy is one of our core values at doc.ai, and privacy is by design in all our products and solutions. We respect the personal information of our users and we understand the choices involved in contributing to better health outcomes via the right vehicle at the right time, such as clinical research or AI trainings. 

To facilitate our twin goals of respecting privacy while learning from data, we have committed ourselves to researching and developing cutting edge privacy-preserving technologies such as edge AI and federated learning. Together these technologies make it possible to train AI models on mobile devices themselves, which allows us in turn to keep an individual’s data on the device, a necessary first step in preserving data privacy.

Consequently we are very excited to share the release of our machine learning framework for on-device inference and training with TensorFlow on Android. We had previously made this framework available for iOS in 2019, and we already supported inference with TensorFlow Lite on Android, but with the release of full TensorFlow support in our open source framework, named Tensor/IO, doc.ai now offers a fully cross-platform solution for prediction, training, and federated learning on any mobile phone.


Edge AI and Federated Learning

Screen Shot 2021-01-07 at 12.53.22 PM


Edge AI is the idea that you can perform AI and specifically machine learning tasks on user devices rather than on server-side systems, for example, on your phone instead of in the cloud. To date Edge AI has been primarily used for predictive purposes and employs frameworks such as TensorFlow Lite, Core ML, and ML Kit.

Federated learning takes the idea of Edge AI a step further and enables algorithms to actually learn from data on user devices without the data ever leaving them. With federated learning we now perform the second fundamental task of machine learning, training, on the device as well, rather than returning that data to the server for learning. In a mobile context that means all user data stays on the phone itself, and the only thing server-side systems see are changes to the model that are a result of training, changes which are later aggregated to produce updated versions of the model.

Together these two technologies, edge A.I and federated learning, constitute the foundation for privacy-preserving algorithmic learning.


The Role of Tensor/IO

While prediction can often be straightforward on Android devices, training is significantly more difficult. Built-in frameworks like ML Kit and mobile machine learning libraries such as TensorFlow Lite can run inference out of the box, but these tools don’t support training. 

To enable training on Android devices doc.ai needed to develop a custom build of the full version of TensorFlow and make it available to the native development system. We took a layered approach to this problem, allowing us to provide simple, high level APIs to the more complex machine learning systems underneath. The end result is that model training can be accomplished in native Android languages like Java and Kotlin in as few as 6-12 lines of code.

The Tensor/IO Stack

Screen Shot 2021-01-07 at 12.53.38 PM

From a more technical standpoint, we began with a set of target models we want to be able to train on the phone, such as MobileNet V2 convolutional models, and we consider our approach successful when we have achieved this goal. Our device stack then started with a custom build of TensorFlow for Android. On top of this we wrote a minimal library bridging Java code into C/C++ code which is capable of calling into TensorFlow. Finally we created our Tensor/IO wrapper to the bridging code which supports high level abstractions, a declarative interface and automatic data transformations, and our cross-platform deployment packaging. At each successive step we tested our example models against the code in that layer of the stack.

At this point Tensor/IO TensorFlow for Android is ready to go, but our own use case required two additional layers to the stack. On top of Tensor/IO we included Tensor/IO Federated, doc.ai’s proprietary federated learning system that involves client-server communication and federated task coordination. Finally, the last layer wrapped the Tensor/IO and Federated code in a React Native module, making the entire stack available to our front-end application teams.

The final product is the complete mobile infrastructure needed to perform prediction, training, and federated learning on Android and iOS devices made available to our application developers in the languages and tools they are already familiar with.


Using Tensor/IO

On January 21 at 9 am PT, we’ll be offering an end-to-end webinar showing you how to train and deploy a model for on-device learning with Tensor/IO. Let’s take a brief look at that process now. We’ll use one of our example models, a MobileNet V2 model trained on the CIFAR 10 dataset, which you can find here: 


After the exploratory analysis and data preparation, the first thing to notice about this example is the use of TensorFlow’s estimator APIs rather than the keras compile, fit, and evaluate APIs:

The estimator APIs are more verbose than the pure keras APIs, but they allow us to explicitly name our training op and to export the model with a training signature rather than a serving (inference) one. This is necessary because the on-device build only has access to the C++ APIs, which means that all training is done via the session APIs with explicit op names and placeholder input dictionaries. If you aren’t familiar with the estimator APIs and the use of input functions and model functions, you can learn more about them here.

In our model_fn notice especially how we name the total_loss and train ops. This will be important later:

Once we’ve trained and evaluated the model using the estimator APIs we can export it for inference or training. In both cases we must set up a serving input receiver function which is ultimately passed to the estimator’s export_saved_model function. This serving input function is responsible for setting up the placeholders which hold the data we’ll be running inference or training on. Notice how we also set up placeholders for the labels in the training version of this function. Otherwise the functions are identical. In particular notice that we name the placeholder layers:

Models are exported in the Saved Model format, which includes the graph definition for the model and a training checkpoint containing the contents of the model’s variables. We’ll need to package this Saved Model data in a Tensor/IO bundle, which is just a folder that includes a model.json file alongside the Saved Model. This model.json file contains everything Tensor/IO needs to use the model, including a description of the model’s input and output layers.

You’ve got to get the model.json file just right for Tensor/IO to be able to use the model. Most importantly you’ll need to match the input and output layer names and shapes to what the model actually provides. For help with this we can use TensorFlow’s saved_model_cli, which is included in the TensorFlow installation. For the prediction model the saved_model_cli shows us:

We’ll want to describe exactly those two input and output layers in our model.json file, which ends up looking like this:

Notice how the names and shapes exactly match the output of the saved_model_cli, which itself matches what was provided in the serving input receiver function. The descriptions also include additional information needed by Tensor/IO to process the model data correctly.

The saved_model_cli output for the training model is quite a bit more complex, but it shows us the three layers we really need for our model, the image, label, and total_loss layers:


We’ll set up our training model.json file the same way using this information but we’ll also add information about the training op, which you’ll recall we explicitly named in our model_fn:


With that model.json file and our Saved Model data packaged into a single folder, we’re now ready to deploy these models to device for inference and training. You can find the Android example for training with this model in the same GitHub repository here:


Let’s take a quick look at the training code in Kotlin:


Almost all of this code is preparatory work. We load the model and an image asset that we’ll use for training, then we set up a batch consisting of that single image asset. Note how the training data has an “image” property and a “label” property. These match the names of the input layers to our model as output by the saved_model_cli and described in the model.json file. 

We then train on the batch for 4 epochs, getting the model’s output each time, in this case the total_loss. Again this “total_loss” exactly matches the name of the output layer which we explicitly named in our model_fn and which is revealed by the saved_model_cli and described in the model.json.

And that’s it! You don’t need to worry about preprocessing the training data or explicitly calling the training operation. Tensor/IO takes care of all of that using the information provided in the model.json file. That’s what we mean when we say that Tensor/IO is a declarative interface to your model. Describe your model in json and Tensor/IO handles all the heavy lifting of preparing data and actually running the model.

That was a lot of information to digest and only a brief overview of the Tensor/IO stack and using Tensor/IO for on-device training. Remember, if you’d like to learn more about building models for on-device training with Tensor/IO, join us for our webinar on Thursday, January 21 at 9 am PT. Register here. Space is limited.

Baseline version


With the release of Tensor/IO TensorFlow for Android, doc.ai now has a fully cross-platform solution for mobile on-device inference and training. With the addition of federated learning, these technologies constitute the first layer in our privacy preserving stack, allowing us to keep medical data on device while still learning from it. At doc.ai data privacy is one of our core values. Edge AI and federated learning allow us to realize this core value and to keep data where it belongs: in your hands.

Learn more about Tensor/IO and contribute at https://github.com/doc-ai/tensorio

Leave a Comment