Training Mobile-Ready models using TensorFlow Lite

Introduction

We can use arcgis.learn to train deep learning models that can be deployed on mobile devices using ArcGIS field apps. This enables AI on the edge and the simplification of field workers' jobs. This is done using TensorFlow which enables us to save trained models to '.tflite' format. A few applications involved in this workflow are

  • Assisted intelligent survey
  • Disconnected Decision support
  • Automatic feature extraction from street view
  • Automated Rapid data collection on the move

image.png

Survey123 application in action

In this guide we are going to use street view images from the ESRI campus. We have labeled a few fire hydrants using labelImg. This is an object detection exercise and requires an object detection model to be trained.

We will be training a SingleShotDetector model in this guide. SingleShotDetector like other common object detection models wil predict the class of the objects along with their location (bounding box) in an image. Learn more about SingleShotDetector here.

image.png

Prerequisites

  • This workflow builds on the basic workflow followed to train deep learning models with arcgis.learn module. It is recommended to first read the basic workflow here if you have not done so previously.
  • To train mobile-ready light models, the TensorFlow library needs to be installed in addition to fastai and PyTorch. Please refer to the section "Using the Deep Learning Frameworks Installer" on this page for detailed documentation on the installation of these dependencies.

About TensorFlow Lite

Deep learning algorithms are generally compute-intensive and specialized hardware like GPUs are always recommended in deep learning workflows. This makes deployment of deep learning models on mobile a challenging task.

TensorFlow Lite is a lightweight solution to this problem. It reduces the overall compute requirement and size of the model making it easy to ship on mobile devices. It also supports hardware acceleration that inturn enables low-latency inferencing on mobile devices, which allows us to build real-time applications on top of it. Learn more about TensorFlow Lite here.

With argis.learn we can export out models to TensorFlow Lite format '.tflite', which can be easily deployed on mobile with ArcGIS field apps. Check the deployment section of this guide for more details on deployment.

Supported Models

Currently, we support three common deep learning workflows for training models that can be used on mobile. Here are the three models along with their application that can be trained for mobile devices:

  • FeatureClassifier - Image Classification
  • SingleShotDetector - Object Detection
  • UnetClassifier - Pixel Classification

Set Environment Variables

This is an important step and has to be run in the first cell of the notebook before any imports are made. By default arcgis.learn only works on a PyTorch backend. Setting this environment variable enables the TensorFlow backend as well, which is required to train a TensorFlow Lite model.

%env ARCGIS_ENABLE_TF_BACKEND=1
env: ARCGIS_ENABLE_TF_BACKEND=1

Imports

from arcgis.learn import prepare_data, SingleShotDetector

Data preparation

data = prepare_data(
    'Esri Campus Street view',
    dataset_type='PASCAL_VOC_rectangles',
    batch_size=16
)

Visualize Training data

data.show_batch()
<Figure size 1152x1152 with 16 Axes>

Train Model

Model instantiation options

  • backbone: Make sure you pick a mobile optimized backbone, currently we only support 'MobileNetV2' backbone.
  • backend: Important !, Make sure you select the backend as 'tensorflow' to enable saving the model to '.tflite' format.
model = SingleShotDetector(data, backbone='MobileNetV2', backend='tensorflow')

We will now find an appropriate learning rate using the {model}.lr_find() method.

model.lr_find()
<Figure size 432x288 with 1 Axes>
0.0004786301

We will now train the model by calling the {model}.fit() method.

model.fit(10, .001)
epochtrain_lossvalid_losstime
083.26815069.82943002:07
164.897545127.27894602:12
253.00910278.50545502:14
347.76968087.28401902:12
442.05009565.23911302:10
538.11028368.72223702:11
634.36035542.28839502:09
732.06900840.12828802:10
829.41141939.24664302:11
928.01295737.34928102:20

Once our model has been fitted for the specified epochs, we can visualize a few sample predictions from the model using the {model}.show_results() method.

Visualizing a few results

model.show_results(thresh=0.5)
<Figure size 576x1440 with 10 Axes>

Save Model to an intermediate format

We can use the {model}.save() method to save the model to 'keras_saved_model' format, this format can be used to save and load intermediate models in the training workflow.

Note: Important! This format is not directly deployable on mobile devices and needs to be exported to mobile ready format as explained in this step.

model.save('10e')
WindowsPath('Esri Campus Street view/models/10e')

Load a model and train it further

If you want to load a model saved to an intermediate format explained here you can use the {model}.load() method.

Note: Important ! Models saved in tflite format explained here cannot be loaded back and trained further as it is an export only format.

# model.load('10e')

Exporting model to mobile ready format

We will now save the model we just trained to a format supported by ESRI Field Apps. Right now the only supported format is 'tflite', and it is required to set the parameter 'framework' to 'tflite' while calling the {model}.save() method.

Note: Important! This is an export only format, and once exported the output cannot be loaded back to the training workflow and fine-tuned again. To save checkpoints or reusable models that can be loaded back, you need to save the model to an intermediate format as explained here.

model.save('10e_tflite', framework="tflite")
WindowsPath('Esri Campus Street view/models/10e_tflite')

Deployment

The tflite model can now be deployed on mobile devices. Survey123 for ArcGIS has an upcoming feature that integrates such tflite models. To learn more on deploying this model in Survey123, join the Early Adopter Community to access the Survey123 private beta.

Your browser is no longer supported. Please upgrade your browser for the best experience. See our browser deprecation post for more details.