17. Exporting precompiled models for TensorFlow Serving

TensorFlow applications compiled for the IPU can be exported to the standard TensorFlow SavedModel format and deployed to a TensorFlow Serving instance. The exported SavedModel contains the executable compiled for the IPU, and a TensorFlow graph with IPU embedded application runtime operations which allow you to run the executable as part of the TensorFlow graph. The exported graph may contain optional preprocessing and postprocessing parts that are executed on the CPU.

The Graphcore TensorFlow API for exporting models for TensorFlow Serving supports three different use cases:

1. Models defined inside a function without using pipelining can be exported using the tensorflow.python.ipu.serving.export_single_step() function.

2. Pipelined models defined as a list of functions can be exported using the tensorflow.python.ipu.serving.export_pipeline() function.

3. Keras models can be exported using tensorflow.python.ipu.serving.export_keras() or the model’s export_for_ipu_serving() method. Both ways are functionally identical and support both pipelined and non-pipelined models.

General notes about using the Graphcore TensorFlow API for exporting models with TensorFlow Serving:

1. Since the exported SavedModel contains custom IPU embedded application runtime operations, it can be used only with the Graphcore distribution of TensorFlow Serving.

2. The exported SavedModel cannot be loaded back into a TensorFlow script and used as a regular model because, in the export stage, the model is compiled into an IPU executable. The exported TensorFlow graph contains only IPU embedded application runtime operations and has no information about specific layers, and so on.

3. TensorFlow and TensorFlow Serving versions must always match. It means that you have to use the same version of TensorFlow Serving as the version of TensorFlow that was used to export the model. Moreover, Poplar versions must also match between the TensorFlow Serving instance and the version of TensorFlow that was used to export the model.

17.1. Exporting non-pipelined models defined inside a function

Exporting the forward pass of a non-pipelined model can be done with the tensorflow.python.ipu.serving.export_single_step() function. A function that defines the forward pass of the model is required as a first argument. Under the hood, export_single_step() wraps that function into a while loop optimized for the IPU, with the iterations parameter denoting the number of loop iterations. You can use this parameter to tweak the model’s latency; its optimal value is use-case specific. Additionally, the function adds the infeed and outfeed queues, so you do not have to take care of it. Then the model is compiled into an executable and included as an asset in the SavedModel stored at the export_dir location. The export_single_step() function adds the possibility of passing the preprocessing_step and postprocessing_step functions which will be included into the SavedModel graph and executed on the CPU on the server side. If all preprocessing and postprocessing operations are available on the IPU, preprocessing_step and postprocessing_step functions should be called inside the predict_step function. Then function bodies will be compiled together with the inference model.

To export such a model, the predict_step function’s input signature has to be defined. This can be accomplished in one of three ways:

• You can decorate the function with the @tf.function decorator which takes the input_signature argument;

• You can pass the predict_step function signature (predict_step_signature) directly to the tensorflow.python.ipu.serving.export_single_step() function;

• You can pass the input dataset (input_dataset) to the tensorflow.python.ipu.serving.export_single_step() function and the exported model’s input signature will be inferred from it.

All of the above methods are functionally equivalent and can be used interchangeably based on what you find more convenient.

17.1.1. Example of exporting non-pipelined model defined inside a function

This example exports a very simple model with embedded IPU program that doubles the input tensor.

import os
import shutil

import numpy as np

import tensorflow as tf

from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving

# Directory where SavedModel will be written.
saved_model_directory = './my_saved_model_ipu/001'
# Directory should be empty or should not exist.
if os.path.exists(saved_model_directory):
  shutil.rmtree(saved_model_directory)


# The function to export.
@tf.function
def predict_step(x):
  # Double the input - replace this with application body.
  result = x * 2
  return result


# Configure the IPU for compilation.
cfg = config.IPUConfig()
cfg.auto_select_ipus = 1
cfg.device_connection.enable_remote_buffers = True
cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
cfg.configure_ipu_system()

input_shape = (4,)
# Prepare the `predict_step` function signature.
predict_step_signature = (tf.TensorSpec(shape=input_shape, dtype=np.float32),)
# Export as a SavedModel.
iterations = 10

runtime_func = serving.export_single_step(predict_step, saved_model_directory,
                                          iterations, predict_step_signature)
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
# This should print the even numbers 0 to 30.
strategy = ipu_strategy.IPUStrategy()
with strategy.scope():
  for i in range(iterations):
    input_data = np.ones(input_shape, dtype=np.float32) * i
    print(runtime_func(input_data))

17.1.2. Example of exporting non-pipelined model defined inside a function with additional preprocessing and postprocessing steps

This example exports a very simple model with an embedded IPU program, which doubles the input tensor. The model also performs a preprocessing step (on the IPU) to compute the absolute value of the input and a postprocessing step (on the IPU) to reduce the output.

import os
import shutil

import numpy as np

import tensorflow as tf

from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving

# Directory where SavedModel will be written.
saved_model_directory = './my_saved_model_ipu/003'
# Directory should be empty or should not exist.
if os.path.exists(saved_model_directory):
  shutil.rmtree(saved_model_directory)


# The preprocessing step is performed fully on the IPU.
def preprocessing_step(x):
  return tf.abs(x)


# The postprocessing step is performed fully on the IPU.
def postprocessing_step(x):
  return tf.reduce_sum(x)


def application_body(x):
  # Double the input - replace this with your application body.
  return x * 2


# The function to export.
@tf.function
def predict_step(x):
  # preprocessing will be compiled and exported together with application body.
  x = preprocessing_step(x)
  x = application_body(x)
  return postprocessing_step(x)


# Configure the IPU for compilation.
cfg = config.IPUConfig()
cfg.auto_select_ipus = 1
cfg.device_connection.enable_remote_buffers = True
cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
cfg.configure_ipu_system()

input_shape = (4,)
# Prepare the `predict_step` function signature.
predict_step_signature = (tf.TensorSpec(shape=input_shape, dtype=np.float32),)
# Export as a SavedModel.
iterations = 10

runtime_func = serving.export_single_step(predict_step, saved_model_directory,
                                          iterations, predict_step_signature)
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
# This should print the even numbers 0 to 30.
strategy = ipu_strategy.IPUStrategy()
with strategy.scope():
  for i in range(iterations):
    input_data = np.ones(input_shape, dtype=np.float32) * (-1.0 * i)
    print(runtime_func(input_data))

This example exports a very simple model with an embedded IPU program, which doubles the input tensor. The model also performs a preprocessing step (on the CPU) to convert string tensors to floats and a postprocessing step (on the CPU) to compute the absolute value of the outputs.

import os
import shutil

import numpy as np

import tensorflow as tf

from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving

# Directory where SavedModel will be written.
saved_model_directory = './my_saved_model_ipu/004'
# Directory should be empty or should not exist.
if os.path.exists(saved_model_directory):
  shutil.rmtree(saved_model_directory)


# The preprocessing step is performed fully on the CPU.
@tf.function
def preprocessing_step(x):
  transform_fn = lambda input: tf.constant(
      1.0) if input == "graphcore" else tf.random.uniform(shape=tuple())

  return tf.stack([transform_fn(elem) for elem in tf.unstack(x)])


# The function to export.
@tf.function
def predict_step(x):
  # preprocessing will be compiled and exported together with application body.
  # Double the input - replace this with your application body.
  return x * 2


# The postprocessing step is performed fully on the CPU.
@tf.function
def postprocessing_step(x):
  return tf.abs(x)


# Configure the IPU for compilation.
cfg = config.IPUConfig()
cfg.auto_select_ipus = 1
cfg.device_connection.enable_remote_buffers = True
cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
cfg.configure_ipu_system()

input_shape = (6,)
# Prepare the `predict_step` function signature.
predict_step_signature = (tf.TensorSpec(shape=input_shape, dtype=np.float32),)
# Prepare the `preprocessing_step` function signature.
preprocessing_step_signature = (tf.TensorSpec(shape=input_shape,
                                              dtype=tf.string),)
# Prepare the `postprocessing_step` function signature.
postprocessing_step_signature = (tf.TensorSpec(shape=input_shape,
                                               dtype=np.float32),)

# Export as a SavedModel.
iterations = 10

runtime_func = serving.export_single_step(
    predict_step,
    saved_model_directory,
    iterations,
    predict_step_signature,
    preprocessing_step=preprocessing_step,
    preprocessing_step_signature=preprocessing_step_signature,
    postprocessing_step=postprocessing_step,
    postprocessing_step_signature=postprocessing_step_signature)
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
strategy = ipu_strategy.IPUStrategy()
with strategy.scope():
  print(
      runtime_func(
          tf.constant(
              ["graphcore", "red", "blue", "yellow", "graphcore", "purple"],
              dtype=tf.string)))
  print(
      runtime_func(
          tf.constant([
              "apple", "banana", "graphcore", "orange", "pineapple",
              "graphcore"
          ],
                      dtype=tf.string)))

17.2. Exporting pipelined models defined as a list of functions

Exporting the forward pass of a pipelined models can be accomplished using tensorflow.python.ipu.serving.export_pipeline() function.

The use of that function is very similar to the creation of a pipeline op using the tensorflow.python.ipu.pipelining_ops.pipeline() function. You have to provide a list of functions that represent the pipeline’s computational stages.

Function tensorflow.python.ipu.serving.export_pipeline() also has an iteration argument. It denotes the number of times each pipeline stage is executed before the pipeline is restarted. Again, you can use it to tweak the model’s latency. This argument is sometimes called steps_per_execution, especially for Keras models.

Similarly to exporting non-pipelined models, to export a pipelined model the signature of the first computational stage has to be known. You can do this in the same three ways as non-pipelined models. It’s worth noting that for the first option—passing the input signature to the @tf.function decorator—you only need to do that for the first computational stage.

17.2.1. Pipeline example

This example exports a simple pipelined IPU program that performs 2x+3 function on the input.

import os
import shutil

import numpy as np
import tensorflow as tf

from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving

# Directory where SavedModel will be written.
saved_model_directory = './my_saved_model_ipu/002'
# Directory should be empty or should not exist.
if os.path.exists(saved_model_directory):
  shutil.rmtree(saved_model_directory)


# The pipeline's stages to export.
def stage1(x):
  # Double the input - replace this with 1st stage body.
  output = x * 2
  return output


def stage2(x):
  # Add 3 to the input - replace this with 2nd stage body.
  output = x + 3
  return output


# Configure the IPU for compilation.
cfg = config.IPUConfig()
cfg.auto_select_ipus = 2
cfg.device_connection.enable_remote_buffers = True
cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
cfg.configure_ipu_system()

input_shape = (4,)
# Prepare the input signature.
predict_step_signature = (tf.TensorSpec(shape=input_shape, dtype=np.float32),)
# Number of times each pipeline stage is executed.
iterations = 10

# Export as a SavedModel.
predict_step = [stage1, stage2]
runtime_func = serving.export_pipeline(
    predict_step,
    saved_model_directory,
    iterations=iterations,
    device_mapping=[0, 1],
    predict_step_signature=predict_step_signature)
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
# This should print numbers from 3 to 33.
strategy = ipu_strategy.IPUStrategy()
with strategy.scope():
  for i in range(iterations):
    input_data = np.ones(input_shape, dtype=np.float32) * i
    print(runtime_func(input_data))

17.2.2. Pipeline example with preprocessing and postprocessing steps

This example exports a simple pipelined IPU program that computes the function 2x+3 on the input. The model includes a preprocessing computational stage which computes the absolute value of the input and an IPU postprocessing step to reduce the output.

import os
import shutil

import numpy as np
import tensorflow as tf

from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving

# Directory where SavedModel will be written.
saved_model_directory = './my_saved_model_ipu/005'
# Directory should be empty or should not exist.
if os.path.exists(saved_model_directory):
  shutil.rmtree(saved_model_directory)


# The preprocessing stage is performed fully on the IPU.
def preprocessing_stage(x):
  return tf.abs(x)


# The pipeline's stages to export.
def stage1(x):
  # Double the input - replace this with 1st stage body.
  output = x * 2
  return output


def stage2(x):
  # Add 3 to the input - replace this with 2nd stage body.
  output = x + 3
  return output


# The postprocessing stage is performed fully on the IPU.
def postprocessing_stage(x):
  return tf.reduce_sum(x)


# Configure the IPU for compilation.
cfg = config.IPUConfig()
cfg.auto_select_ipus = 4
cfg.device_connection.enable_remote_buffers = True
cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
cfg.configure_ipu_system()

input_shape = (4,)
# Prepare the input signature.
predict_step_signature = (tf.TensorSpec(shape=input_shape, dtype=np.float32),)
# Number of times each pipeline stage is executed.
iterations = 8

# Export as a SavedModel.
predict_step = [preprocessing_stage, stage1, stage2, postprocessing_stage]
runtime_func = serving.export_pipeline(
    predict_step,
    saved_model_directory,
    iterations=iterations,
    device_mapping=[0, 1, 2, 3],
    predict_step_signature=predict_step_signature)
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
# This should print numbers from 3 to 33.
strategy = ipu_strategy.IPUStrategy()
with strategy.scope():
  for i in range(iterations):
    input_data = np.ones(input_shape, dtype=np.float32) * (-1.0 * i)
    print(runtime_func(input_data))

This example exports a simple pipelined IPU program that computes the function 2x+3 on the input tensor. The model also performs a preprocessing step (on the CPU) to convert string tensors to floats and a postprocessing step (on the CPU) to compute the absolute value of the outputs.

import os
import shutil

import numpy as np
import tensorflow as tf

from tensorflow.python.ipu import config
from tensorflow.python.ipu import ipu_strategy
from tensorflow.python.ipu import serving

# Directory where SavedModel will be written.
saved_model_directory = './my_saved_model_ipu/006'
# Directory should be empty or should not exist.
if os.path.exists(saved_model_directory):
  shutil.rmtree(saved_model_directory)


# The preprocessing stage is performed fully on the CPU.
@tf.function
def preprocessing_step(x):
  transform_fn = lambda input: tf.constant(
      1.0) if input == "graphcore" else tf.random.uniform(shape=tuple())

  return tf.stack([transform_fn(elem) for elem in tf.unstack(x)])


# The pipeline's stages to export.
def stage1(x):
  # Double the input - replace this with 1st stage body.
  output = x * 2
  return output


def stage2(x):
  # Add 3 to the input - replace this with 2nd stage body.
  output = x + 3
  return output


# The postprocessing step is performed fully on the CPU.
@tf.function
def postprocessing_step(x):
  return tf.abs(x)


# Configure the IPU for compilation.
cfg = config.IPUConfig()
cfg.auto_select_ipus = 2
cfg.device_connection.enable_remote_buffers = True
cfg.device_connection.type = config.DeviceConnectionType.ON_DEMAND
cfg.configure_ipu_system()

input_shape = (6,)
# Prepare the input signature.
predict_step_signature = (tf.TensorSpec(shape=input_shape, dtype=np.float32),)
# Prepare the `preprocessing_step` function signature.
preprocessing_step_signature = (tf.TensorSpec(shape=input_shape,
                                              dtype=tf.string),)
# Prepare the `postprocessing_step` function signature.
postprocessing_step_signature = (tf.TensorSpec(shape=input_shape,
                                               dtype=np.float32),)

# Number of times each pipeline stage is executed.
iterations = 10

# Export as a SavedModel.
predict_step = [stage1, stage2]
runtime_func = serving.export_pipeline(
    predict_step,
    saved_model_directory,
    iterations=iterations,
    device_mapping=[0, 1],
    predict_step_signature=predict_step_signature,
    preprocessing_step=preprocessing_step,
    preprocessing_step_signature=preprocessing_step_signature,
    postprocessing_step=postprocessing_step,
    postprocessing_step_signature=postprocessing_step_signature)
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
# This should print numbers from 3 to 33.
strategy = ipu_strategy.IPUStrategy()
with strategy.scope():
  print(
      runtime_func(
          tf.constant(
              ["graphcore", "red", "blue", "yellow", "graphcore", "purple"],
              dtype=tf.string)))
  print(
      runtime_func(
          tf.constant([
              "apple", "banana", "graphcore", "orange", "pineapple",
              "graphcore"
          ],
                      dtype=tf.string)))

17.3. Exporting Keras models

There are two ways of exporting Keras models for TensorFlow Serving, independent of whether they’re pipelined or not. Keras models can be exported using the tensorflow.python.ipu.serving.export_keras() function or the model’s export_for_ipu_serving() method.

See the Section 19, Keras with IPUs section for details and examples of exporting precompiled Keras models for TensorFlow Serving.

17.4. Running the model in TensorFlow Serving

To test the exported SavedModel you can just start a TensorFlow Serving instance and point it to the model’s location. Graphcore’s distribution of TensorFlow Serving can be run directly in the host system:

$ tensorflow_model_server --rest_api_port=8501 --model_name=my_model \
      --model_base_path="$(pwd)/my_saved_model_ipu"

And then you can start sending inference requests, for example:

$ curl -d '{"instances": [1.0, 2.0, 5.0, 7.0]}'   \
    -X POST http://localhost:8501/v1/models/my_model:predict

Graphcore does not distribute the TensorFlow Serving API package. If you want to use it, you need to install it from the official distribution using pip.