Skip to content

TensorFlow Backend

Philip Dow edited this page May 14, 2020 · 9 revisions

The following has been tested with TensorFlow 1.13, 1.14, and 1.15 using the Estimator APIs. We have also tested some models built in TF 2.0. See tensorio/examples for a collection of notebooks and ios sample code demonstrating how to use the models built with them.

Inference

The TensorFlow backend supports both inference and training. Instructions for preparing a model for training are below. Preparing models for inference is straightforward.

Prep

Activate your tensorflow virtual environment or make sure you have the tensorflow package installed, and at the command line confirm that the saved_model_cli is available:

$ which saved_model_cli
/path/to/tensorflow/bin/saved_model_cli

Export

Export your model to the saved_model format using any of a number of the built in tensorflow APIs, as you normally would to package your model for inference. We use the estimator API with a serving input receiver function:

def save_model(model_dir, output_dir, dims):
  estimator = tf.estimator.Estimator(
    model_fn=model_fn, 
    model_dir=model_dir)
  estimator.export_saved_model(
    output_dir, 
    lambda:serving_input_receiver_fn(),
    as_text=False)

This should produce a directory with the following structure, which is effectively the inference graph with a checkpoint:

exported_model/
  saved_model.pb
  variables/
    variables.data-00000-of-00001
    variables.index

Saved Model CLI

Get information about your model's input and output layers using tensorflow's saved model cli. We need the name and the shape of our input and output nodes for Tensor/IO's model.json file:

$ saved_model_cli show --all --dir exported_model

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['image'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 128, 128, 3)
        name: image:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['class'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 1)
        name: ToInt32_1:0
    outputs['probability'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: sigmoid:0
  Method name is: tensorflow/supervised/training

In this case I have an input layer named image (the trailing ':0' can be disregarded unless you have more than one layer with this name), and two output layers. The output layer I'm interested in is sigmoid. Including the batch dimension, which is relevant to Tensor/IO's TensorFlow backend, the shapes are [-1,128,128,3] and [-1,1], respectively.

Model.json

We'll use this information about layer names and shapes to prepare our model.json file. The relevant fields are found in the "inputs" and "outputs" keys. Note that the input corresponds to the image layer in our model along with its shape and that the output corresponds to the sigmoid layer and its shape. We specify the RGB format and pixel normalization in [0,1] because that is how this model has been set up. Your values may be different.

Finally, note in the "model" field that the "modes" key is ["predict"].

{
  "name": "Image Prediction Model",
  "details": "Based on MobileNetV2 architecture",
  "id": "whatever-you-want",
  "version": "1",
  "author": "doc.ai",
  "license": "Apache License. Version 2.0 http://www.apache.org/licenses/LICENSE-2.0",
  "model": {
    "file": "exported_model",
    "quantized": false,
    "type": "image.classification.example",
    "backend": "tensorflow",
    "modes": ["predict"]
  },
  "inputs": [
    {
      "name": "image",
      "type": "image",
      "shape": [-1,128,128,3],
      "format": "RGB",
      "normalize": {
        "standard": "[0,1]"
      }
    }
  ],
  "outputs": [
    {
      "name": "sigmoid",
      "type": "array",
      "shape": [-1,1]
    }
  ]
}

Packaging the Model

First note in the model.json file that the "model.file" field must be the name of the folder that contains the saved_model export. I typically use "predict" here but we named our folder "exported_model" in the above example, so'll we continue to use that.

With you model.json ready and your exported_model folder in hand, you are now ready to package this model for use with Tensor/IO. Simply create a new directory with the .tiobundle extension and place the model.json and exported_model folder in it. The directory structure will look like:

my_model.tiobundle/
  model.json
  exported_model/
    saved_model.pb
    variables/
      variables.data-00000-of-00001
      variables.index

Add that tiobundle to your project. You may now load and use that model from your app.

For more information about packaging models for deployment, see Packaging-Models

Training

Preparing models for training on device with the TensorFlow backend requires a little more preparation. You will need to make a number of changes to your model_fn and then you will have to export the model for training using an experimental mode and a modified serving input receiver function.

Prep

Activate your tensorflow virtual environment or make sure you have the tensorflow package installed, and at the command line confirm that the saved_model_cli is available:

$ which saved_model_cli
/path/to/tensorflow/bin/saved_model_cli

The model_fn

You must make a number of changes to your model_fn to export a model for training.

  1. Acquire the model's labels from the feature paramater
  2. Explicitly name the call to your optimizer's minimization operation
  3. Remove unsupported ops

Additional Input Placeholders

When exporting a model for training, what were previously the model's outputs (labels) become training inputs. This is explicit with the Session API but less so with the Estimator API. Consider your model_fn definition:

def model_fn(features, labels, mode, params):
  # ...

When exporting for training, the features parameter will be set as expected, and the mode will be set to tf.estimator.ModeKeys.TRAIN, but the labels parameter will be unfilled and have a value of None. We have to check for that in the model_fn code and get the placeholders for the labels from the features parameter instead. We'll make sure the features parameter has a labels key with modifications to the serving input receiver function below.

For now, add this near the top of your model function before labels is used:

if labels is None: # during training export
  labels = features['labels']

Name the Minimization Op

Because Tensor/IO uses the C++ Session API it must know the name of the training op to call in order to actually execute a round of training.

Find the call to your optimizer's minimize function that yields the training op, or whatever else produces your training op, and explicitly name it:

optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
train_op = optimizer.minimize(
  loss=loss,
  global_step=tf.train.get_global_step(),
  name="train")

Here we've named it "train", and we'll refer to that name later on in the model.json file.

Remove Unsupported Ops

You may find by a process of trial and error that your graph contains ops that are not supported by Tensor/IO's mobile build of TensorFlow. We've tried to package that build with all the ops required for running common edge models, but some ops simply can't be supported, for example, the summary ops. Find those ops and comment them out when you export for training. Or fork on a flag and only include them when you are not exporting for training.

# unsupported ops on mobile build so comment them out

# tf.summary.scalar('average_loss', loss)
# tf.summary.scalar('accuracy', accuracy[1])

Export

You are now ready to export your model for training! Export your model to the saved_model format using any of a number of the built in tensorflow APIs, as you normally would to package your model for inference. We use the estimator API with a serving input receiver function, and our example has a single image input.

The important difference here is that we must 1) set up our serving input receiver function so that it also has a placeholder for the labels and 2) add the experimental_mode parameter to the estimator.export_saved_model call.

Normally my serving input receiver function might look like this:

def serving_input_prediction_receiver_fn(params):
  """
  The serving input receiver function used for prediction only takes 
  the image as an input
  """

  image_shape  = [None, params['target_dim'], params['target_dim'], 3]

  inputs = {
    'image': tf.placeholder(tf.float32, shape=image_shape, name='image'),
  }

  return tf.estimator.export.ServingInputReceiver(inputs, inputs)

To export a model for training, I need to include a placeholder for the labels. Remember that we had to pull the labels from the feature parameter to the model_fn above because the labels parameter is not set when exporting for training. Our serving input receiver function is effectively setting up the value to that parameter.

Now my serving input receiver function looks like:

def serving_input_training_receiver_fn(params):
  """
  The serving input receiver function used for training takes the image
  as well as the label for input
  """

  image_shape  = [None, params['target_dim'], params['target_dim'], 3]
  labels_shape = [None,1]

  inputs = {
    'image': tf.placeholder(tf.float32, shape=image_shape, name='image'),
    'labels': tf.placeholder(tf.int32, shape=labels_shape, name='labels')
  }

  return tf.estimator.export.ServingInputReceiver(inputs, inputs)

Notice that I explicitly name the labels placeholder. We'll refer to that name later on in the model.json file. Do this for any of your model's outputs that are now being treated like inputs for training.

Export the model with that new serving input receiver function, and set the experimental_mode parameter in the estimator.export_saved_model call to tf.estimator.ModeKeys.TRAIN:

def save_model(model_dir, output_dir, dims):
  # set up params to the serving input receiver function with dims
  params = # ...
  
  estimator = tf.estimator.Estimator(
    model_fn=model_fn, 
    model_dir=model_dir)
  estimator.export_saved_model(
    output_dir, 
    lambda:serving_input_training_receiver_fn(params),
    as_text=False,
    experimental_mode=tf.estimator.ModeKeys.TRAIN)

This may strike you as odd, but it is how it's done. This should produce a directory with the following structure, which is effectively the training graph with a checkpoint:

exported_model/
  saved_model.pb
  variables/
    variables.data-00000-of-00001
    variables.index

Saved Model CLI

Get information about your model's input and output layers using tensorflow's saved model cli. We need the name and the shape of our input and output nodes for Tensor/IO's model.json file:

$ saved_model_cli show --all --dir exported_model

MetaGraphDef with tag-set: 'train' contains the following SignatureDefs:

signature_def['train']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['image'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 128, 128, 3)
        name: image:0
    inputs['labels'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 1)
        name: labels:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['loss'] tensor_info:
        dtype: DT_FLOAT
        shape: ()
        name: sigmoid_cross_entropy_loss/value:0
  Method name is: tensorflow/supervised/training

Notice that we have a labels key to the inputs! That's thanks to our modified serving input receiver function. We also have a model output in the signature def, but it doesn't have anything to do with the model's labels. Instead it is the loss value. When you export a model for training, tensorflow considers the model's output to be its loss. The name of this particular node is sigmoid_cross_entropy_loss/value which is taken from the loss function we used in our example model:

loss = tf.losses.sigmoid_cross_entropy(
  multi_class_labels=labels,
  logits=logits
)

We will also want to report this loss during mobile training, so note the name of that output for use in the model.json file.

As for inputs, we have an image input with a shape and label with a shape. In this case the image layer is named image (the trailing ':0' can be disregarded unless you have more than one layer with this name) and its shape is [-1,128,128,3]. Don't take the key to the inputs dictionary as the name, take the layer's name instead. Our labels input is named labels and its shape is [-1,1]. We'll refer to these values in the model.json file.

Model.json

We'll use the information about layer names and shapes, the loss layer name, and the optimization op name to prepare our model.json file. The relevant fields are found in the "inputs", "outputs", and "train" keys. Note that we have two inputs corresponds to the image layer in our model along with its shape and the labels placeholder and its shape, that the output corresponds to the loss layer and its shape, and that we also including the training op by name. For the image input we specify the RGB format and pixel normalization in [0,1] because that is how this model has been set up. Your values may be different.

Finally, note in the "model" field that the "modes" key is ["train"].

{
  "name": "Image Training Model",
  "details": "Based on MobileNetV2 architecture",
  "id": "whatever-you-want",
  "version": "1",
  "author": "doc.ai",
  "license": "Apache License. Version 2.0 http://www.apache.org/licenses/LICENSE-2.0",
  "model": {
    "file": "exported_model",
    "quantized": false,
    "type": "image.classification.example",
    "backend": "tensorflow",
    "modes": ["train"]
  },
  "inputs": [
    {
      "name": "image",
      "type": "image",
      "shape": [-1,128,128,3],
      "format": "RGB",
      "normalize": { "standard": "[0,1]" }
    },
    {
      "name": "labels",
      "type": "array",
      "dtype": "int32",
      "shape": [-1,1]
    }
  ],
  "outputs": [
    {
      "name": "sigmoid_cross_entropy_loss/value",
      "type": "array",
      "shape": [1]
    }
  ],
  "train": {
    "ops": [
      "train"
    ]
  }
}

Packaging the Model

Packaging a model for training is the same as packaging one for inference.

First note in the model.json file that the "model.file" field must be the name of the folder that contains the saved_model export. For trainable models I typically use "train" here but we named our folder "exported_model" in the above example, so'll we continue to use that.

With you model.json ready and your exported_model folder in hand, you are now ready to package this model for use with Tensor/IO. Simply create a new directory with the .tiobundle extension and place the model.json and exported_model folder in it. The directory structure will look like:

my_model.tiobundle/
  model.json
  exported_model/
    saved_model.pb
    variables/
      variables.data-00000-of-00001
      variables.index

Add that tiobundle to your project. You may now load and train with that model from your app!

For more information about packaging models for deployment, see Packaging-Models

~

Last updated 5/14/20.

Clone this wiki locally