TensorFlow Mobile: Training and Deploying a Neural Network

Smart Assistants, fancy image filters in Snapchat and apps like Prisma all have one thing in common—they are powered by Machine Learning. The use of Machine Learning in mobile apps is growing and new mobile apps are developed with Machine Learning based services as business models. In this blog series we want to give you hands-on advice on how you can train and deploy a convolutional neural network for image classification to a mobile app using the popular machine learning framework TensorFlow Mobile.

Our task will be to classify images of houseplants which we have collected ourselves. You don’t have to go and snap pictures of plants, however, because our approach is generic and can be used for training and deploying a convolutional neural network for image classification, independent of their subject. If you’d also like to go with houseplants, however, we have written an image crawler to save you the manual labor. You’ll find the instructions here.

For your convenience we have published a repository containing all necessary files and source code used in this tutorial.

As a concrete implementation of a convolutional neural network we’ll use one of the MobileNets, a class of efficient convolutional neural networks for mobile and embedded vision applications. These are already implemented in one of the high-level APIs of TensorFlow which is called TF-Slim. You can find the TF-Slim models in the model repository of TensorFlow. In this blog series we will use TF-Slim for the training of the MobileNet.

For the deployment of neural networks to a mobile device there are currently two solutions:

• Tensor Flow Mobile: TensorFlow was designed from the ground up to be a good deep learning solution for mobile platforms such as Android and iOS. TensorFlow Mobile represents the mobile version of the framework which you can use in your mobile apps. It also contains multiple guides and scripts for the deployment of a model into a mobile app.
• TensorFlow Lite: This is an evolution of TensorFlow Mobile. In most cases, apps developed with TensorFlow Lite will have a smaller binary size, fewer dependencies, and better performance. Currently TensorFlow Lite is in developer preview, so not all use cases are covered yet and it only supports a limited set of operators, so not all models will work on it by default.

In this blog series we will use TensorFlow Mobile because TensorFlow Lite is in developer preview and TensorFlow Mobile has a greater feature set. As mentioned before, we will use images of houseplants as our dataset. In total there are 9364 images across 26 classes available.

The Setup

First off we need to install TensorFlow. The easiest way is to follow the official installation guide which works for different platforms and operation systems. Just use the current version of TensorFlow.

To get started we clone the TensorFlow model repository:  git clone https://github.com/tensorflow/models.git and switch to the TF-Slim models directory with cd models/research/slim. There we find the following tree structure:

Next we need to define our dataset, creating a python file as description in the dataset directory, where the other dataset descriptions are arranged. With  cd datasets we switch to the directory and with  touch hp_plants.py we create the required python file, to which we add the following code:

You can use every dataset you want, you just have to change the name of the python file and of the dataset variables inside the python file. The dataset variables are marked in the code above. Next we need to make a reference to our dataset in the dataset factory. Open  dataset_factory.py in the datasets folder and add the name of the dataset to the datasets map.

To convert our image data to an appropriate binary file format (TFRecord) we use a script provided by Kwotsin in a Github repository. Clone the repository and copy  create_tfrecord.py and  dataset_utils.py to the slim folder. To create the TFRecord files run the following command in your terminal.

With the dataset_dir parameter we define where our dataset is stored and with the  tfrecord_filename parameter we define the pattern of the TFRecord files. This pattern must match with the pattern we defined in our dataset description and the dataset factory. In the last parameter you define which size of the dataset should be used for validation. This step will create TFRecord files for training and validation. With the current setting of the  validation_size parameter, 80 % of the data will be used for training and 20 % for validation. The results can be viewed in the dataset folder of the houseplants.

The execution of the script has created 2 TFRecord files for execution and 2 for validation. Also a file with the labels was created, which contains the 26 class names of the dataset. The images folder contains the images of the houseplants in particular folders. For every class there is a folder with the class name as folder name and the images of the class inside of this particular folder.

Training

Until now we have done general preparation and pre-processing. In the next steps we will set up our training, using  Transfer Learning. In practice an entire convolutional neural network is rarely trained from scratch, because it is rare to have a dataset of sufficient size. With Transfer Learning however we can train a convolutional neural network with a dataset of a small size, because we are using pre-trained weights of the convolutional neural network. We just have to fine-tune it on our dataset.

In the model repository of TensorFlow you can download multiple pre-trained weights of several different convolutional neural networks trained on ImageNet data. As mentioned above we are using a MobileNet in this blogpost series, whose pre-trained weights we have to download. We can find them in the MobileNet v1 description where we have to download MobileNet_v1_1.0_224. Copy the downloaded .tgz file to the slim folder, create a subfolder with the name  mobilenet_v1_1.0_224 and extract it with  tar xf mobilenet_v1_1.0_224.tgz -C ./mobilenet_v1_1.0_224 to the subfolder. In the subfolder you can see multiple files.

The weights for transfer learning are stored in the .ckpt files.

TensorFlow uses a dataflow graph to represent computations in terms of the dependencies between individual operations. Dataflow is a common programming model for parallel computing where the nodes represent units of computation and the edges represent the data consumed or produced, which also applies to neural networks in TensorFlow. In the subfolder we have a whole graph of the MobileNet, which is stored in the .pb file provided. We’ll later need this  for to provision our model for mobile.

To start our training we need to run  train_image_classifier.py with some arguments. We recommend training on a GPU to speed up the process considerably. Run this command in your terminal to start training on the GPU:

If you don’t have a GPU available, you have to use the following command. With the argument clone_on_cpu=True all computations will be executed on CPU.

With the argument  train_dir we specify the destination of the TFRecord files we created beforehand. As a further argument we have  dataset_name to select the dataset we want to use for training. Here we choose the houseplants dataset previously specified in the datasets description. The  dataset_split argument specifies which TFRecord files are used for training, so we select those from above. In the next argument  model_name we specify which model we want to train. Here we have choose the MobileNet as model, setting our input images size to 224x224x3 with the argument  train_image_size . With the argument  checkpoint_path we refer to the downloaded checkpoint of the MobileNet while also enabling Transfer Learning as training method. The argument  max_number_step defines the number of training steps, which we set too 30,000.

In the end we need to exclude some weights of the checkpoint, because the model was pre-trained on the ImageNet datatset. The weights of the fully connected layer, which does the classification in the end, are trained on the ImageNet dataset with 1000 classes and our houseplants dataset has 26 classes. So we can not use the pre-trained weights for this layer and have to train it completely from scratch. During the training you see the loss converge with the output looking like this:

The model will be automatically saved in the specified train directory. After that we need to evaluate the trained model using the provided python script  eval_image_classifier.py. Run the script with the following command:

It’s very important to refer to the right checkpoint for evaluation. This we can specify with the argument  checkpoint_path where we refer to the checkpoint of the 30000 training steps. The output looks like this:

We can see that our MobileNet trained with Transfer Learning in 30000 steps achieves an accuracy of  about 75 % and a top-5-recall of 92 %. Thats quite good for such a short training. To boost the performance we could raise the number of training steps but this may lead to overfitting. Another approach is called Hyperparameter Optimization which automatically searches for good parameters like learning rate or regularization strength for our neural network. 

Mobile Deployment

After we have finished training and evaluation of our MobileNet, we now can start with the preparation of the mobile deployment. For this we first need to create an inference graph of our MobileNet, which represents the whole MobileNet and is used to map our trained weights with the correct graph of the MobileNet. To create the inference graph we need to run this command:

You can now find the correct graph representation of the MobileNet in the slim folder as .pb file with the name  inference_graph_mobilenet.pb. Now we need to freeze our trained MobileNet, mapping the trained weights with the correct graph representation of the MobileNet:

As you can see our previously generated inference graph is used as input for the freezing. Also we are using the trained weights from our latest checkpoint. With the argument  output_graph we specify the output name of the frozen graph. Furthermore we need to provide the argument  output_node with the right output node name. The information about input and output nodes of the MobileNet can be found in the previously downloaded files at  ./mobilenet_v1_1.0_224/mobilenet_v1_1.0_224_info.txt.

As a last step we need to optimize our graph for mobile. This step will reduce the binary size of the graph by removing unnecessary operations for classificaton and rounding the provided weights of the model. Rounding the provided weights will lead to a small accuracy loss but will improve the classification duration of the model greatly which very is important for mobile devices. To optimize our graph we need to run the the following command:

As you can see the input for the optimization is our previously frozen graph which will be optimized for mobile and which is saved as .pb file at  opt_frozen_mobilenet.pb.

Now we have a fully functional and mobile optimized graph which we can deploy to our Android or iOS app—we’ll show you how in future articles!

Read on to learn how to integrate the graph with your Android app!