19. Keras with IPUs

The Graphcore implementation of Keras includes support for the IPU. Keras model creation is no different than what you would use if you were training on other devices. To target the Poplar XLA device, Keras model creation must be inside the strategy.scope of an IPUStrategy.

For a more practical walkthrough, see this tutorial about using Keras on the IPU from the Graphcore tutorials repository.

19.1. Single IPU models

You can train, evaluate or run inference on single-IPU models through the Keras APIs as you would with other accelerators, as long as you create the model inside the scope of an IPUStrategy:

import tensorflow as tf
from tensorflow.python import ipu

import keras
from keras.datasets import mnist

# Configure the IPU device.
config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()


# Create a simple model.
def create_model():
  return keras.Sequential([
      keras.layers.Flatten(),
      keras.layers.Dense(256, activation='relu'),
      keras.layers.Dense(128, activation='relu'),
      keras.layers.Dense(10)
  ])


# Create a dataset for the model.
def create_dataset():
  (x_train, y_train), (_, _) = mnist.load_data()
  x_train = x_train / 255.0

  train_ds = tf.data.Dataset.from_tensor_slices(
      (x_train, y_train)).shuffle(10000).batch(32, drop_remainder=True)
  train_ds = train_ds.map(lambda d, l:
                          (tf.cast(d, tf.float32), tf.cast(l, tf.int32)))

  return train_ds.repeat().prefetch(16)


dataset = create_dataset()

# Create a strategy for execution on the IPU.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Create a Keras model inside the strategy.
  model = create_model()

  # Compile the model for training.
  model.compile(
      loss=keras.losses.SparseCategoricalCrossentropy(),
      optimizer='rmsprop',
      metrics=["accuracy"],
  )

  model.fit(dataset, epochs=2, steps_per_epoch=100)

19.2. Using steps_per_execution

To reduce Python overhead and maximize the performance of your model, pass the steps_per_execution argument to the compile method. This argument sets the number of batches processed sequentially by one replica in a single execution which can greatly improve performance because any overhead between steps is removed, thus increasing IPU utilization.

Ideally, steps_per_execution is equal to the number of steps your model needs to run per replica in order to complete one epoch. Note that it is not possible to fetch intermediate results when steps_per_execution is specified. Model weights are read on the Python host after all steps are executed on the IPU. If you need to access model weights during an epoch (for example for saving a checkpoint), you must set steps_per_execution accordingly.

Note

In order to achieve best performance, steps_per_execution needs to be set before using fit(), evaluate() and predict(), even if no training is performed.

Note

To achieve best performance when using pipelining, steps_per_execution should be set to a significantly larger value than the number of pipeline stages. If steps_per_execution is too small for pipelining to work, your model will fail to compile and the minimum value of steps_per_execution will be reported in the error message.

See the documentation for the compile method for full details.

The example below highlights the usage of steps_per_execution:

import tensorflow as tf
from tensorflow.python import ipu

import keras
from keras.datasets import mnist

# Configure the IPU device.
config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()


# Create a simple model.
def create_model():
  return keras.Sequential([
      keras.layers.Flatten(),
      keras.layers.Dense(256, activation='relu'),
      keras.layers.Dense(128, activation='relu'),
      keras.layers.Dense(10)
  ])


# Create a dataset for the model.
def create_dataset():
  (x_train, y_train), (_, _) = mnist.load_data()
  x_train = x_train / 255.0

  train_ds = tf.data.Dataset.from_tensor_slices(
      (x_train, y_train)).shuffle(10000).batch(32, drop_remainder=True)
  train_ds = train_ds.map(lambda d, l:
                          (tf.cast(d, tf.float32), tf.cast(l, tf.int32)))

  return train_ds.prefetch(16)


dataset = create_dataset()

# Create a strategy for execution on the IPU.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Create a Keras model inside the strategy.
  model = create_model()

  # Compile the model for training.
  model.compile(
      loss=keras.losses.SparseCategoricalCrossentropy(),
      optimizer='rmsprop',
      metrics=["accuracy"],
      # Anything between 2 and the length of the dataset would work,
      # but the greater `steps_per_execution` the greater the
      # performance gains.
      steps_per_execution=dataset.cardinality(),
  )

  model.fit(dataset, epochs=2)

19.3. Gradient accumulation

When training, gradient accumulation allows us to simulate bigger batch sizes. This is achieved by accumulating the gradients across multiple batches together then performing the weight update.

For example, if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica to 4 then this simulates an input batch of size 64.

Gradient accumulation can be easily enabled for Keras models created inside of an IPUStrategy by calling the following methods:

Functional model	set_gradient_accumulation_options()
Sequential model	set_gradient_accumulation_options()
Model subclass	set_gradient_accumulation_options()

This method can be used to configure gradient accumulation parameters in models that do not use pipelining, particularly gradient_accumulation_steps_per_replica and gradient_accumulation_reduction_method. See the respective API documentation for more details on these arguments.

For pipelined models, these arguments can be passed to the set_pipelining_options() methods. See Section 19.4.4, Pipelining options for more details

Note

A step commonly refers to the forward and backward passes, followed by a weight update. On the IPU, when gradient accumulation is used, a step refers to the forward and backward passes on a micro batch, but not including the corresponding weight update (the weight update is only performed every gradient_accumulation_steps_per_replica steps). The number of weight update steps per execution is given by the steps_per_execution value the model was compiled with, divided by gradient_accumulation_steps_per_replica. An execution, which is a compiled Poplar program, must have an integer number of weight update steps, such that all the accumulated gradients are applied. Therefore, steps_per_execution must be an integer multiple of gradient_accumulation_steps_per_replica.

Note

The steps per epoch value (steps_per_epoch) applies per replica. steps_per_epoch needs to be set when using a dataset of infinite cardinality. An epoch consists of one or more executions. Therefore, if set, steps_per_epoch must be an integer multiple of steps_per_execution.

Note

Not all operations are compatible with gradient accumulation.

The example below highlights the usage of set_gradient_accumulation_options:

import tensorflow as tf
from tensorflow.python import ipu

import keras
from keras.datasets import mnist

# Configure the IPU device.
config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()


# Create a simple model.
def create_model():
  return keras.Sequential([
      keras.layers.Flatten(),
      keras.layers.Dense(256, activation='relu'),
      keras.layers.Dense(128, activation='relu'),
      keras.layers.Dense(10)
  ])


# Create a dataset for the model.
def create_dataset():
  (x_train, y_train), (_, _) = mnist.load_data()
  x_train = x_train / 255.0

  train_ds = tf.data.Dataset.from_tensor_slices(
      (x_train, y_train)).shuffle(10000).batch(32, drop_remainder=True)
  train_ds = train_ds.map(lambda d, l:
                          (tf.cast(d, tf.float32), tf.cast(l, tf.int32)))

  return train_ds.prefetch(16)


dataset = create_dataset()

# Create a strategy for execution on the IPU.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Create a Keras model inside the strategy.
  model = create_model()

  # `steps_per_execution` must be divisible by `gradient_accumulation_steps_per_replica`.
  # Say we want to accumulate 10 steps before doing a weight update, then we would end up
  # with the following values.
  gradient_accumulation_steps_per_replica = 10
  number_of_accumulated_steps = dataset.cardinality(
  ) // gradient_accumulation_steps_per_replica

  # In order to get the proper `steps_per_execution` value, we have to multiply
  # `number_of_accumulated_steps` with `gradient_accumulation_steps_per_replica`.
  steps_per_execution = number_of_accumulated_steps * \
                        gradient_accumulation_steps_per_replica

  # Now we need to truncate the dataset so Keras will not try to take more data
  # from the dataset than is available.
  dataset = dataset.take(steps_per_execution)

  # Compile the model for training.
  model.compile(
      loss=keras.losses.SparseCategoricalCrossentropy(),
      optimizer='rmsprop',
      metrics=["accuracy"],
      steps_per_execution=steps_per_execution,
  )

  model.set_gradient_accumulation_options(
      gradient_accumulation_steps_per_replica=10)

  model.fit(dataset, epochs=2)

19.4. Model parallelism

The models described so far occupy a single IPU device, however some models might require the model layers to be split across multiple IPU devices to achieve high compute efficiency.

One method to achieve model parallelism is called pipelining, where the model layers are assigned to pipeline stages. Each pipeline stage can be assigned to a different device and different devices can execute in parallel.

By default, these pipeline stages will be executed using the grouped schedule (Fig. 19.1), where the forward and backward stages are grouped together on each IPU. All IPUs alternate between executing a forward pass and then a backward pass.

Fig. 19.1 Grouped pipeline

Two other schedules are available and can be configured as shown in Section 19.4.4, Pipelining options. When using the interleaved schedule (Fig. 19.2) the forward and backward passes are interleaved (which requires less memory but is likely to be slower). The sequential schedule (Fig. 19.3) executes one stage at a time and may be useful when debugging your model.

Fig. 19.2 Interleaved pipeline

Fig. 19.3 Sequential pipeline

A detailed explanation of pipelining can be found in the technical note on Model parallelism with TensorFlow: sharding and pipelining.

Recomputation of activations can be enabled to reduce memory usage. See Section 6.6, Recomputation for further details on when and how to use it.

The method to pipeline your model depends on whether your model is a Sequential model, a Functional model, or is subclassed from the Model class.

19.4.1. Sequential model

Enabling pipelining for a Sequential model requires assigning each layer in the model to a pipeline stage, by calling set_pipeline_stage_assignment().

If your model does not contain any nested Keras models, a simple overload of set_pipeline_stage_assignment() can be used which accepts a list of integers. The integers in the list correspond to the layers in the model, each specifying which pipeline stage the corresponding layer is assigned to.

For example, a simple four layer Sequential model could be assigned to two different pipeline stages as follows:

  model = keras.Sequential([
      keras.layers.Dense(8),  # Pipeline stage 0.
      keras.layers.Dense(16),  # Pipeline stage 0.
      keras.layers.Dense(16),  # Pipeline stage 1.
      keras.layers.Dense(1),  # Pipeline stage 1.
  ])

  model.set_pipeline_stage_assignment([0, 0, 1, 1])

You can confirm which layers are assigned to which stages using the print_pipeline_stage_assignment_summary() method of the model.

Pipelining a model containing nested models

If your Sequential model contains nested Keras models, you can use get_pipeline_stage_assignment(). This returns a list of pipeline stage assignment objects, corresponding to the layers in the model, which you can inspect and modify.

Each layer in the model has an associated SequentialLayerPipelineStageAssignment object, which indicates what pipeline stage that layer is assigned to.

Each nested Keras model in the model has an associated SequentialNestedModelPipelineStageAssignment containing a list of pipeline stage assignments which you can inspect and modify.

Once you are done modifying the stage assignments, you can pass them to set_pipeline_stage_assignment() to set them on the model. Pipeline stage assignments for nested Keras models will NOT be recursively set on each nested Keras model. The assignments are all stored on the root model (the model on which set_pipeline_stage_assignment() was called) and they are only used when calling fit(), predict(), or evaluate() on that model.

19.4.2. Functional model

There are two ways to enable IPU pipelining for a Functional model (an instance of keras.Model) depending on if you’re pipelining a model you are writing yourself or an existing model.

Pipelining a model you are writing yourself

To pipeline a Functional model you are writing yourself, each layer call must happen within the scope of an keras.ipu.PipelineStage context.

For example, a simple four layer Functional model could be assigned to two different pipeline stages as follows:

  input_layer = keras.layers.Input((28, 28))

  with keras.ipu.PipelineStage(0):
    x = keras.layers.Dense(8)(input_layer)
    x = keras.layers.Dense(16)(x)

  with keras.ipu.PipelineStage(1):
    x = keras.layers.Dense(16)(x)
    x = keras.layers.Dense(1)(x)

  model = keras.Model(inputs=input_layer, outputs=x)

Note

Layers constructed within a PipelineStage context will have that pipeline stage assigned to all invocations of the layer. These assignments are overridden if the layer calls happen within a different PipelineStage context.

Pipelining an existing functional model

To pipeline an existing Functional model, you can use get_pipeline_stage_assignment(). This returns a list of pipeline stage assignment objects, corresponding to each invocation in the model which you can inspect and modify. Note that the list is in post-order, which means the assignments are returned in the order they will be executed.

Each layer invocation in the model has an associated FunctionalLayerPipelineStageAssignment object, which indicates what pipeline stage that invocation is assigned to.

Each nested Keras model invocation in the model has an associated FunctionalNestedModelPipelineStageAssignment containing a list of pipeline stage assignments which you can inspect and modify.

Once you are done modifying the stage assignments, you can pass them to set_pipeline_stage_assignment() to set them on the model. Pipeline stage assignments for nested Keras models will NOT be recursively set on each nested Keras model. The assignments are all stored on the root model (the model on which set_pipeline_stage_assignment() was called) and they are only used when calling fit(), predict(), or evaluate() on that model.

For example, a naive way of pipelining ResNet50 would be to assign everything up until the “conv4_block2_add” layer invocation to the first stage, then everything else to the second stage, as follows:

strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():

  model = resnet.ResNet50(weights='imagenet')

  # Get the individual assignments - note that they are returned in post-order.
  assignments = model.get_pipeline_stage_assignment()

  # Iterate over them and set their pipeline stages.
  stage_id = 0
  for assignment in assignments:
    assignment.pipeline_stage = stage_id
    # Split the model on the `conv4_block2_add` layer.
    if assignment.layer.name.startswith("conv4_block2_add"):
      stage_id = 1

  # Set the assignments to the model.
  model.set_pipeline_stage_assignment(assignments)

  model.print_pipeline_stage_assignment_summary()

Note

You can use print_pipeline_stage_assignment_summary() to print the pipeline stage assignments of the model’s layer invocations.

Note

This method of assigning pipeline stages can also be used with Functional models you are writing yourself, as well as with Sequential models and Model subclasses using the SequentialExtension and ModelExtension equivalents.

19.4.3. Model subclass

Model subclasses are subclasses of keras.Model, which override the call method. There are two ways to enable IPU pipelining for an instance of a Model subclass, depending on if you’re pipelining a model you are writing yourself or an existing model. These are very similar to the methods available for Functional models.

Note that the following Model methods cannot be overridden:

• train_step

• make_train_function

• test_step

• make_test_function

• predict_step

• make_predict_function

Pipelining a model you are writing yourself

To pipeline a Model subclass you are writing yourself, each layer call must happen within the scope of an keras.ipu.PipelineStage context.

For example, a simple four layer Model subclass could be assigned to four different pipeline stages as follows:

class MyModel(keras.Model):
  def __init__(self):
    super().__init__(self)
    self.dense_layer_1 = keras.layers.Dense(8)
    self.dense_layer_2 = keras.layers.Dense(8)
    self.concat_layer = keras.layers.Concatenate()
    self.dense_layer_3 = keras.layers.Dense(1)

  def call(self, inputs):
    # Invoke layers inside PipelineStage scopes to assign the layer invocations
    # to the specified pipeline stage.
    with keras.ipu.PipelineStage(0):
      x = self.dense_layer_1(inputs)
    with keras.ipu.PipelineStage(1):
      x1 = self.dense_layer_2(x)
      x2 = self.dense_layer_2(x)
    with keras.ipu.PipelineStage(2):
      x1 = self.dense_layer_2(x1)
      x2 = self.dense_layer_2(x2)
      x = self.concat_layer([x1, x2])
    with keras.ipu.PipelineStage(3):
      x = self.dense_layer_3(x)

    return x

Note

Layers constructed within a PipelineStage context will have that pipeline stage assigned to all invocations of the layer. These assignments are overridden if the layer calls happen within a different PipelineStage context.

Pipelining an existing model

To pipeline an existing Model subclass, you must use get_pipeline_stage_assignment(). This returns a list of pipeline stage assignment objects, corresponding to each invocation in the model which you can inspect and modify. Note that the list is in post-order, which means the assignments are returned in the order they will be executed.

Each layer invocation in the model has an associated ModelLayerPipelineStageAssignment object, which indicates what pipeline stage that invocation is assigned to.

Each nested Keras model invocation in the model has an associated NestedModelPipelineStageAssignment containing a list of pipeline stage assignments which you can inspect and modify.

Once you are done modifying the stage assignments, you can pass them to set_pipeline_stage_assignment() to set them on the model. Pipeline stage assignments for nested Keras models will NOT be recursively set on each nested Keras model. The assignments are all stored on the root model (the model on which set_pipeline_stage_assignment() was called) and they are only used when calling fit(), predict(), or evaluate() on that model.

Before you can get or set pipeline stage assignments, you must first call keras.Model.build() on your model, specifying the input shapes. This traces the model’s call function using the shapes specified. The resulting graph is what will be used for pipelined execution. You can update the graph by calling build again, though this will invalidate existing pipeline stage assignments if the structure of the updated graph is different.

Note

If you need to specify input dtypes when calling keras.Model.build(), you can pass in keras.Input objects instead of plain shapes.

For example, an existing Model subclass with four layers, could be assigned to four different pipeline stages as follows:

  model = ExistingModel()

  # Call build to trace the graph generated by the call function.
  # This step is required before getting or setting pipeline stage assignments.
  model.build((28, 28))

  # Get a blank set of pipeline stage assignments.
  assignments = model.get_pipeline_stage_assignment()

  # Modify the assignments by setting pipline stages.
  for assignment in assignments:
    if assignment.layer == model.dense_layer_1:
      assignment.pipeline_stage = 0
    elif assignment.layer == model.dense_layer_2 and assignment.node_index < 2:
      assignment.pipeline_stage = 1
    elif assignment.layer == model.dense_layer_2 and assignment.node_index < 4:
      assignment.pipeline_stage = 2
    elif assignment.layer == model.concat_layer:
      assignment.pipeline_stage = 2
    elif assignment.layer == model.dense_layer_3:
      assignment.pipeline_stage = 3

  # Apply the modified assignments back to the model.
  model.set_pipeline_stage_assignment(assignments)

Note

You can use print_pipeline_stage_assignment_summary() to print the pipeline stage assignments of the model’s layer invocations.

Note

This method of assigning pipeline stages can also be used with Model subclasses you are writing yourself, as well as with Functional and Sequential models using the SequentialExtension and FunctionalExtension equivalents.

19.4.4. Pipelining options

Pipelining options can be set with the following methods:

Functional model	set_pipelining_options()
Sequential model	set_pipelining_options()
Model subclass	set_pipelining_options()

See the respective API documentation for more details.

Gradient accumulation is always used when training a pipelined model (unless using the Sequential schedule). This means that you must set the option gradient_accumulation_steps_per_replica using this API when using the Grouped or Interleaved schedule. It is optional when using the Sequential schedule.

The API documentation for set_pipelining_options explains that the additional keyword arguments (pipelining_kwargs) will be forwarded to the tensorflow.python.ipu.pipelining_ops.pipeline() operator (which is used internally - see Section 19.11, Implementation details). Refer to the API documentation for pipeline() for details about these arguments.

The code sample below illustrates how options can be set with the set_pipelining_options API.

  model.set_pipelining_options(
      gradient_accumulation_steps_per_replica=16,
      pipeline_schedule=ipu.ops.pipelining_ops.PipelineSchedule.Interleaved)

19.5. Automatic data parallelism

IPU TensorFlow supports automatic data parallelism when multiple IPU devices are configured with the system. Automatic data parallelism is achieved by model replication across available IPU devices. The number of times the model is replicated is called the replication factor; higher replication factors allow higher data throughput.

When using replication, gradients are reduced across replicas during training, which has implications for gradient accumulation. For a non replicated model, the effective batch size is the product of the dataset batch size and the number of gradient accumulation steps. In the case of a replication factor greater than one, the effective batch size is additionally scaled by the replication factor according to the following formula:

effective_batch_size = dataset_batch_size * gradient_accumulation_steps_per_replica * num_replicas

Metrics can also be reduced across replicas. This behaviour must be configured by calling the following methods on your model, and specifying a value for replicated_metric_reduction_method:

Functional model	set_replication_options()
Sequential model	set_replication_options()
Model subclass	set_replication_options()

See the respective API documentation for more details.

19.6. Asynchronous callbacks

IPU TensorFlow supports the use of Callback objects with the Keras APIs, however there is an important difference to note when specifying steps_per_execution. In IPU TensorFlow, if steps_per_execution is specified for your model, then per-batch callback functions will only be invoked every steps_per_execution steps, which can have the effect of delaying access to results.

However, IPU TensorFlow also supports asynchronous callbacks by providing a polling mechanism which allows results to be accessed at the earliest possible instance. Asynchronous callbacks can be enabled by passing True to the following methods:

Functional model	set_asynchronous_callbacks()
Sequential model	set_asynchronous_callbacks()
Model subclass	set_asynchronous_callbacks()

See the respective API documentation for more details.

19.7. Configuring Infeeds and Outfeed

Keras models created inside of an IPUStrategy scope automatically create IPUInfeedQueue and OutfeedQueue data queues for efficiently feeding data to and from the IPU devices when using fit(), evaluate() and predict().

Instances of IPUInfeedQueue and OutfeedQueue can be created with optional arguments which can affect performance of the model.

Use the following methods to configure the IPUInfeedQueue for your Keras model:

Functional model	set_infeed_queue_options()
Sequential model	set_infeed_queue_options()
Model subclass	set_infeed_queue_options()

Use the following methods to configure the OutfeedQueue for your Keras model:

Functional	set_outfeed_queue_options()
Sequential	set_outfeed_queue_options()
Model subclass	set_outfeed_queue_options()

For example the prefetch_depth parameter of the OutfeedQueue and the buffer_depth parameter of the OutfeedQueue can be configured as follows:

from tensorflow.python import ipu

import keras

# Configure the IPU device.
config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()


# Create a simple model.
def create_model():
  return keras.Sequential([
      keras.layers.Flatten(),
      keras.layers.Dense(256, activation='relu'),
      keras.layers.Dense(128, activation='relu'),
      keras.layers.Dense(10)
  ])


# Create a strategy for execution on the IPU.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():

  model = create_model()

  # Set the infeed and outfeed options.
  model.set_infeed_queue_options(prefetch_depth=2)
  model.set_outfeed_queue_options(buffer_depth=2)

19.8. Saving and loading Keras models

Saving and loading a Keras model must be done within the IPUStrategy scope in order to save/load IPU-specific information.

When saving and loading Model subclasses, make sure to save and restore class members, such as layers, via the config. This can be done by overriding the get_config and from_config methods. Re-creating members from scratch can cause errors, as the original members may be restored as part of the IPU-specific internal state.

Note

The arguments pipelining_kwargs from set_pipelining_options() and gradient_accumulation_optimizer_kwargs from set_gradient_accumulation_options() are not serializable, which means that when the model is being saved, their values are not saved. When restoring/loading a model, call set_pipelining_options() or set_gradient_accumulation_options() again.

19.9. Exporting precompiled Keras models for TensorFlow Serving

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. This takes only three arguments: the model to export, a directory where the SavedModel will be stored and, optionally, a batch size value. The other way uses the model’s export_for_ipu_serving() method which takes only the path to the SavedModel directory and, optionally, a batch size value.

It’s important to note that before exporting the model you must build it, providing the input shapes to the model’s build() method. Similarly to exporting non-Keras models, you can set the iteration parameter by calling the model’s compile() method with steps_per_execution argument. The meaning of that parameter is analogous to that of non-Keras models, both non-pipelined and pipelined ones. In both cases you can use it to tweak the inference latency.

The export_for_ipu_serving() method 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. If all preprocessing and postprocessing operations are available on the IPU, preprocessing_step and postprocessing_step functions should be called inside the Keras model. Then function bodies will be compiled together with the inference model.

Exported models contain Poplar programs compiled for specific batch size value. Because of that, you must always provide the batch size value to be used by the exported model. You can achieve it in two ways:

• passing the batch_size argument explicitly to the export function, or

• setting the batch size value during model creation and leaving the default value of the batch_size argument.

19.9.1. Non-pipelined Keras model example

This example creates a simple non-pipelined Keras model that adds two inputs together. After that, the model is exported for TensorFlow Serving.

import os
import shutil

import numpy as np
from tensorflow.python import ipu

import keras

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

batch_size = 1
input_shape = (batch_size, 4)
# Number of IPU-optimized loop's iterations.
iterations = 16

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

# Always create Keras models inside an IPU strategy.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Always set `batch_size` if model has explicit input layers.
  input1 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_1")
  input2 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_2")
  x = keras.layers.Add()([input1, input2])
  model = keras.Model(inputs=[input1, input2], outputs=x)

  model.build([input_shape, input_shape])
  # Call compile to set the number of iterations of the inference loop.
  # It can be used to tweak the inference latency.
  model.compile(steps_per_execution=iterations)

# Export as a SavedModel.
runtime_func = model.export_for_ipu_serving(saved_model_directory)
# Alternatively: `runtime_func = serving.export_keras(model, saved_model_directory)`
print(f"SavedModel written to {saved_model_directory}")

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

19.9.2. Non-pipelined Keras model example with additional preprocessing and postprocessing steps

This example exports a very simple Keras model with an embedded IPU program that adds two inputs together. The model also performs a preprocessing step (on the IPU) to compute the absolute value of the input tensors 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 import ipu

import keras

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

batch_size = 1
input_shape = (batch_size, 4)
# Number of IPU-optimized loop's iterations.
iterations = 16

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


# The preprocessing step is performed fully on the IPU.
def preprocessing_step(lhs_input, rhs_input):
  abs_layer = keras.layers.Lambda(tf.abs)
  return abs_layer(lhs_input), abs_layer(rhs_input)


# The postprocessing step is performed fully on the IPU.
def postprocessing(model_result):
  reduce_layer = keras.layers.Lambda(tf.reduce_sum)
  return reduce_layer(model_result)


# Always create Keras models inside an IPU strategy.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Always set `batch_size` if model has explicit input layers.
  input1 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_1")
  input2 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_2")

  x = keras.layers.Add()(preprocessing_step(input1, input2))
  output = postprocessing(x)

  model = keras.Model(inputs=[input1, input2], outputs=output)

  model.build([input_shape, input_shape])
  # Call compile to set the number of iterations of the inference loop.
  # It can be used to tweak the inference latency.
  model.compile(steps_per_execution=iterations)

# Export as a SavedModel.
runtime_func = model.export_for_ipu_serving(saved_model_directory)
# Alternatively: `runtime_func = serving.export_keras(model, saved_model_directory)`
print(f"SavedModel written to {saved_model_directory}")

# You can test the exported executable using returned `runtime_func`.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  for i in range(iterations):
    input1_data = np.ones(input_shape, dtype=np.float32) * i
    input2_data = np.ones(input_shape, dtype=np.float32) * 2
    print(runtime_func(input1_data, input2_data))

This example exports a very simple Keras 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 import ipu

import keras

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

batch_size = 1
input_shape = (batch_size, 6)
# Number of IPU-optimized iterations.
iterations = 16

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

# Prepare the `preprocessing_step` function signature.
preprocessing_step_signature = (tf.TensorSpec(shape=input_shape,
                                              dtype=tf.string),
                                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),)


# The preprocessing step is performed fully on the CPU.
@tf.function(input_signature=preprocessing_step_signature)
def preprocessing_step(lhs_input, rhs_input):
  transform_fn = lambda input: tf.constant(
      1.0) if input == "graphcore" else tf.random.uniform(shape=tuple(),
                                                          dtype=np.float32)
  transform_string = lambda input: tf.stack([
      tf.stack([transform_fn(elem) for elem in tf.unstack(rank1)])
      for rank1 in tf.unstack(input)
  ])
  return transform_string(lhs_input), transform_string(rhs_input)


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


# Always create Keras models inside an IPU strategy.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Always set `batch_size` if model has explicit input layers.
  input1 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_1")
  input2 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_2")

  x = keras.layers.Add()([input1, input2])

  model = keras.Model(inputs=[input1, input2], outputs=x)

  model.build([input_shape, input_shape])
  # Call `compile` to set the number of iterations of the inference loop.
  # It can be used to tweak the inference latency.
  model.compile(steps_per_execution=iterations)

# Export as a SavedModel.
runtime_func = model.export_for_ipu_serving(
    saved_model_directory,
    preprocessing_step=preprocessing_step,
    postprocessing_step=postprocessing_step)
# Alternatively: `runtime_func = serving.export_keras(
#   model,
#   saved_model_directory,
#   preprocessing_step=preprocessing_step,
#   postprocessing_step=postprocessing_step)`
print(f"SavedModel written to {saved_model_directory}")

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

19.9.3. Pipelined Keras model example

This example creates a simple pipelined Keras model that adds two inputs together in the first pipeline stage and later multiplies the result of the addition operation with the second input in the second pipeline stage. After that, the model is exported for TensorFlow Serving.

Note that building, compiling and exporting look exactly the same for pipelined and non-pipelined models.

import os
import shutil

import numpy as np
from tensorflow.python import ipu

import keras

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

batch_size = 1
input_shape = (batch_size, 4)
iterations = 16

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

# Always create Keras models inside an IPU strategy.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Always set `batch_size` if model has explicit input layers.
  input1 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_1")
  input2 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_2")

  with keras.ipu.PipelineStage(0):
    x = keras.layers.Multiply()([input1, input2])

  with keras.ipu.PipelineStage(1):
    x = keras.layers.Add()([x, input2])

  model = keras.Model(inputs=[input1, input2], outputs=x)
  model.set_pipelining_options(device_mapping=[0, 1])

  model.build([input_shape, input_shape])
  # Call compile to set the number of times each pipeline stage is executed.
  # It can be used to minimize the latency a bit.
  model.compile(steps_per_execution=iterations)

# Export as a SavedModel.
runtime_func = model.export_for_ipu_serving(saved_model_directory)
# Alternatively: `runtime_func = serving.export_keras(model, saved_model_directory)`
print("SavedModel written to", saved_model_directory)

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

19.9.4. Pipelined Keras model example with additional preprocessing and postprocessing steps

This example creates a simple pipelined Keras model that adds two inputs together in the first computational pipeline stage of the model and later multiplies the result of the addition operation with the second input in the next pipeline stage. The model also performs a preprocessing stage (on the IPU) to compute the absolute value of the input and a postprocessing stage (on the IPU) to reduce the output.

import os
import shutil

import numpy as np
import tensorflow as tf
from tensorflow.python import ipu

import keras

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

batch_size = 1
input_shape = (batch_size, 4)
iterations = 16

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


# The preprocessing step is performed fully on the IPU.
def preprocessing_step(lhs_input, rhs_input):
  abs_layer = keras.layers.Lambda(tf.abs)
  return abs_layer(lhs_input), abs_layer(rhs_input)


# The postprocessing step is performed fully on the IPU.
def postprocessing(model_result):
  reduce_layer = keras.layers.Lambda(tf.reduce_sum)
  return reduce_layer(model_result)


# Always create Keras models inside an IPU strategy.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Always set `batch_size` if model has explicit input layers.
  input1_ph = keras.layers.Input(shape=input_shape[1:],
                                 batch_size=batch_size,
                                 name="input_1")
  input2_ph = keras.layers.Input(shape=input_shape[1:],
                                 batch_size=batch_size,
                                 name="input_2")
  with keras.ipu.PipelineStage(0):
    input1, input2 = preprocessing_step(input1_ph, input1_ph)

  with keras.ipu.PipelineStage(1):
    x = keras.layers.Multiply()([input1, input2])

  with keras.ipu.PipelineStage(2):
    x = keras.layers.Add()([x, input2])

  with keras.ipu.PipelineStage(3):
    x = postprocessing(x)

  model = keras.Model(inputs=[input1_ph, input2_ph], outputs=x)
  model.set_pipelining_options(device_mapping=[0, 1, 2, 3])

  model.build([input_shape, input_shape])
  # Call compile to set the number of times each pipeline stage is executed.
  # It can be used to minimize the latency a bit.
  model.compile(steps_per_execution=iterations)

# Export as a SavedModel.
runtime_func = model.export_for_ipu_serving(saved_model_directory)
# Alternatively: `runtime_func = serving.export_keras(model, saved_model_directory)`
print("SavedModel written to", saved_model_directory)

# You can test the exported executable using returned runtime_func
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  for i in range(iterations):
    input1_data = np.ones(input_shape, dtype=np.float32) * i
    input2_data = np.ones(input_shape, dtype=np.float32) * 2
    print(runtime_func(input1_data, input2_data))

This example creates a simple pipelined Keras model that adds two inputs together in the first pipeline stage and later multiplies the result of the addition operation with the second input in the second pipeline stage. 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 import ipu

import keras

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

batch_size = 1
input_shape = (batch_size, 6)
iterations = 16

# Prepare the `preprocessing_step` function signature.
preprocessing_step_signature = (tf.TensorSpec(shape=input_shape,
                                              dtype=tf.string),
                                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),)


# The preprocessing step is performed fully on the CPU.
@tf.function(input_signature=preprocessing_step_signature)
def preprocessing_step(lhs_input, rhs_input):
  transform_fn = lambda input: tf.constant(
      1.0) if input == "graphcore" else tf.random.uniform(shape=tuple(),
                                                          dtype=np.float32)

  transform_string = lambda input: tf.stack([
      tf.stack([transform_fn(elem) for elem in tf.unstack(rank1)])
      for rank1 in tf.unstack(input)
  ])
  return transform_string(lhs_input), transform_string(rhs_input)


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


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

# Always create Keras models inside an IPU strategy.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Always set `batch_size` if model has explicit input layers.
  input1 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_1")
  input2 = keras.layers.Input(shape=input_shape[1:],
                              batch_size=batch_size,
                              name="input_2")

  with keras.ipu.PipelineStage(0):
    x = keras.layers.Multiply()([input1, input2])

  with keras.ipu.PipelineStage(1):
    x = keras.layers.Add()([x, input2])

  model = keras.Model(inputs=[input1, input2], outputs=x)
  model.set_pipelining_options(device_mapping=[0, 1])

  model.build([input_shape, input_shape])
  # Call compile to set the number of times each pipeline stage is executed.
  # It can be used to minimize the latency a bit.
  model.compile(steps_per_execution=iterations)

# Export as a SavedModel.
runtime_func = model.export_for_ipu_serving(
    saved_model_directory,
    preprocessing_step=preprocessing_step,
    postprocessing_step=postprocessing_step)
# Alternatively: `runtime_func = serving.export_keras(
#   model,
#   saved_model_directory,
#   preprocessing_step=preprocessing_step,
#   postprocessing_step=postprocessing_step)`

print("SavedModel written to", saved_model_directory)

# You can test the exported executable using returned runtime_func
# This should print the even numbers 2 to 32.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  input1_data = tf.constant(
      ["graphcore", "red", "blue", "yellow", "graphcore", "purple"],
      shape=input_shape,
      dtype=tf.string)
  input2_data = tf.constant(
      ["apple", "banana", "graphcore", "orange", "pineapple", "graphcore"],
      shape=input_shape,
      dtype=tf.string)
  print(runtime_func(input1_data, input2_data))

19.10. IPU-specific Keras layers and optimizers

The ipu_tensorflow_addons.keras.layers namespace contains IPU-specific implementations of standard Keras layers and optimizers. More information, including details of every layer and optimizer in this namespace and a code example showing how to use it can be found in Section 20, IPU TensorFlow Addons.

19.11. Implementation details

When instantiating a standard TensorFlow Keras model inside the scope of an IPUStrategy instance, it is dynamically injected with additional, IPU-specific, functions. This is done through the relevant IPU Keras extension classes:

Functional model	FunctionalExtension()
Sequential model	SequentialExtension()
Model subclass	ModelExtension()

19.12. Automatic loss scaling

When training deep learning models, the magnitude of computed gradients is typically significantly smaller than their corresponding weights and activations. Whilst rarely causing training issues for models using float32 precision weights, models using reduced precision formats like float16 are left vulnerable to numerical underflow and vanishing gradients.

Loss scaling aims to combat vanishing gradients by multiplying the loss value at the end of the forward pass by some loss scaling factor \({\alpha}\). Consequently, gradients obtained through back-propagation are also scaled by that constant. The aim is to shift the gradient distribution across the dynamic range, so that underflow and overflow are prevented (as much as possible) in float16. Note that the gradients need to be scaled down by the inverse of \({\alpha}\) before being consumed by the optimizer for the weight update.

Automatic loss scaling (ALS), as included with IPU Keras, eliminates the need for the manual selection of an appropriate value for \({\alpha}\). The ALS approach observes the gradients with respect to both weights and activations. These observations are used to generate histograms that inform the adjustment of the loss scaling factor, with the aim of preventing excessive underflow or overflow in gradient distribution. The proportion of FP16 gradients whose absolute values exceed histogram_bin_edge is recorded. If this proportion is below ratio_threshold then \({\alpha}\) is scaled up by increase_factor. Otherwise, it is scaled down by 1 / increase_factor. These updates to \({\alpha}\) occur with a user-specified update_frequency. For more details and visual examples, see our blogpost about ALS.

Warning

The automatic loss scaling feature in the IPU version of Keras is experimental and may lead to unexpected results. So far it has been validated for the SGD optimizer without momentum and with constant learning rate. You can see an example of it in our CNNs application.

Note

update_frequency uses the number of calls to apply_gradients to determine how often to update \({\alpha}\). Features affecting the frequency of these calls, like replication, pipelining or gradient accumulation, might therefore require corresponding modifications to update_frequency for ALS to exhibit the desired behavior.

The continual updating of \({\alpha}\) has the added benefit, as compared to a static scaling factor, of allowing the optimizer to adapt to changes in the distribution of magnitudes of gradients during training.

In IPU Keras, ALS can be added to any OptimizerV2-derived Keras optimizer through the ALSOptimizer wrapper. The example below illustrates how the ALSOptimizer wrapper can be used to add ALS functionality to a standard SGD optimizer and train a model.

# Simple training example using ALS.

import numpy as np

from tensorflow.python import ipu

import keras
from keras.ipu.optimizers import ALSOptimizer
from keras.optimizer_v2.gradient_descent import SGD

config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()

strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  input_layer = keras.layers.Input(32)
  x = keras.layers.Dense(16)(input_layer)
  x = keras.layers.Dense(8)(x)

  m = keras.Model(input_layer, x)

  opt = SGD(0.01)
  opt = ALSOptimizer(opt)

  m.compile(loss='mse', optimizer=opt)

  in_data = np.ones((128, 32), dtype=np.float16)
  targets = 2.0 * np.ones((128, 8), dtype=np.float16)

  m.fit(in_data, targets, epochs=10)

While the example above uses a keras.Model for simplicity, the ALSOptimizer wrapper can also be used within custom TensorFlow training loops, as shown in the example below.

# Simple training example using ALS without keras.Model.

import numpy as np

import tensorflow as tf
from tensorflow.python import ipu
from tensorflow.python.ipu import loops

import keras
from keras.ipu.optimizers import ALSOptimizer
from keras.optimizer_v2.gradient_descent import SGD

config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()

strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  opt = SGD(0.01)
  opt = ALSOptimizer(opt)

  layer_0 = keras.layers.Dense(16)
  layer_1 = keras.layers.Dense(8)
  mse = keras.losses.MeanSquaredError(
      reduction=keras.losses.losses_utils.ReductionV2.SUM)

  @tf.function(jit_compile=True)
  def f(x, t, _):
    y = layer_1(layer_0(x))
    l = mse(y_true=t, y_pred=y)

    v = layer_0.trainable_variables + layer_1.trainable_variables
    g = opt.get_gradients(l, v)
    opt.apply_gradients(zip(g, v))

    return x, t, l

  @tf.function(jit_compile=True)
  def training_loop(x, t):
    _, _, l = loops.repeat(10, f, inputs=[x, t, 0.0])
    return l

  in_data = np.ones((128, 32), dtype=np.float16)
  targets = 2.0 * np.ones((128, 8), dtype=np.float16)

  strategy.run(training_loop, args=[in_data, targets])

The following example shows how ALS can be combined with gradient accumulation (outlined in Section 19.3, Gradient accumulation) in order to simulate larger batch sizes.

# Simple training example using ALS with gradient accumulation.

import numpy as np

from tensorflow.python import ipu

import keras
from keras.ipu.optimizers import ALSOptimizer
from keras.optimizer_v2.gradient_descent import SGD

config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()

strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  input_layer = keras.layers.Input(32)
  x = keras.layers.Dense(16)(input_layer)
  x = keras.layers.Dense(8)(x)

  m = keras.Model(input_layer, x)

  opt = SGD(0.01)
  opt = ALSOptimizer(opt)

  m.compile(loss='mse', optimizer=opt, steps_per_execution=8)
  m.set_gradient_accumulation_options(
      gradient_accumulation_steps_per_replica=4)

  in_data = np.ones((128, 32), dtype=np.float16)
  targets = 2.0 * np.ones((128, 8), dtype=np.float16)

  m.fit(in_data, targets, epochs=10, steps_per_epoch=8)

The combination of ALS and gradient accumulation does not require the use of keras.Model, however. The example below illustrates how ALSGradientAccumulationOptimizer can be used to provide any OptimizerV2-derived Keras optimizer with both ALS and gradient accumulation.

Note

In order to use ALSGradientAccumulationOptimizer, the target optimizer must first be wrapped with ALSOptimizer.

# Simple training example using ALS and gradient accumulation
# without keras.Model.

import numpy as np

import tensorflow as tf
from tensorflow.python import ipu
from tensorflow.python.ipu import loops

import keras
from keras.ipu.optimizers import ALSOptimizer
from keras.ipu.optimizers import ALSGradientAccumulationOptimizer
from keras.optimizer_v2.gradient_descent import SGD

config = ipu.config.IPUConfig()
config.auto_select_ipus = 1
config.configure_ipu_system()

strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  opt = SGD(0.01)
  opt = ALSOptimizer(opt)
  opt = ALSGradientAccumulationOptimizer(opt, num_mini_batches=2)

  layer_0 = keras.layers.Dense(16)
  layer_1 = keras.layers.Dense(8)
  mse = keras.losses.MeanSquaredError(
      reduction=keras.losses.losses_utils.ReductionV2.SUM)

  @tf.function(jit_compile=True)
  def f(x, t, _):
    y = layer_1(layer_0(x))
    l = mse(y_true=t, y_pred=y)

    v = layer_0.trainable_variables + layer_1.trainable_variables
    g = opt.get_gradients(l, v)
    opt.apply_gradients(zip(g, v))

    return x, t, l

  @tf.function(jit_compile=True)
  def training_loop(x, t):
    _, _, l = loops.repeat(10, f, inputs=[x, t, 0.0])
    return l

  in_data = np.ones((128, 32), dtype=np.float16)
  targets = 2.0 * np.ones((128, 8), dtype=np.float16)

  strategy.run(training_loop, args=[in_data, targets])

While it is ultimately the gradients with respect to weights and biases which are used to update the parameters of a model, gradients with respect to layer activations are also computed during each backward pass. It is possible that these gradients might saturate, thus losing information, without this being detected in upstream parameter gradients.

By wrapping a Keras layer in the CaptureActivationGradients wrapper, these intermediate activation gradients can be recorded and added to the statistics used to determine whether to increase or decrease \({\alpha}\). The example below shows how the activation gradients of select layers can be incorporated into an ALSOptimizer. In order to reduce the computational cost of ALS, the update_frequency is increased to 5 (from its default value of 1), and the accumulate_statistics_over_update_period flag is set to False. With this setting, the ratio of gradients exceeding histogram_bin_edge are recorded only every fifth batch, rather than being based on the average ratio over the last 5 batches.

# MNIST training with ALS.
# Uses gradients w.r.t activations and only collects statistics
# every 5 iterations.

import tensorflow as tf
from tensorflow.python import ipu

import keras
from keras.datasets import mnist
from keras.ipu.optimizers import ALSOptimizer
from keras.ipu.layers import CaptureActivationGradients
from keras.optimizer_v2 import gradient_descent

# Configure the IPU device.
config = ipu.config.IPUConfig()
config.auto_select_ipus = 2
config.configure_ipu_system()


# Create a simple model.
def create_model():
  input_layer = keras.layers.Input((28, 28, 1))

  with keras.ipu.PipelineStage(0):
    x = CaptureActivationGradients(
        keras.layers.Conv2D(32,
                            dtype=tf.float16,
                            kernel_size=(3, 3),
                            activation="relu"))(input_layer)
    x = keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
    x = CaptureActivationGradients(
        keras.layers.Conv2D(64,
                            dtype=tf.float16,
                            kernel_size=(3, 3),
                            activation="relu"))(x)

  with keras.ipu.PipelineStage(1):
    x = keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
    x = keras.layers.Flatten()(x)
    x = keras.layers.Dropout(0.5)(x)
    x = CaptureActivationGradients(keras.layers.Dense(10,
                                                      activation="softmax"))(x)

  m = keras.Model(inputs=input_layer, outputs=x)
  m.set_pipelining_options(gradient_accumulation_steps_per_replica=8)

  return m


# Create a dataset for the model.
def create_dataset():
  (x_train, y_train), (_, _) = mnist.load_data()
  x_train = x_train / 255.0
  x_train = tf.expand_dims(x_train, -1)

  train_ds = tf.data.Dataset.from_tensor_slices(
      (x_train, y_train)).shuffle(10000).batch(32, drop_remainder=True)
  train_ds = train_ds.map(lambda d, l:
                          (tf.cast(d, tf.float32), tf.cast(l, tf.int32)))

  return train_ds.repeat().prefetch(16)


# Create the optimizer.
def create_optimizer():
  o = gradient_descent.SGD(0.01)
  return ALSOptimizer(o,
                      update_frequency=5,
                      accumulate_statistics_over_update_period=False)


dataset = create_dataset()

# Create a strategy for execution on the IPU.
strategy = ipu.ipu_strategy.IPUStrategy()
with strategy.scope():
  # Create a Keras model inside the strategy.
  model = create_model()

  # Create the optimizer and ALS wrapper.
  opt = create_optimizer()

  # Compile the model for training.
  model.compile(loss=keras.losses.SparseCategoricalCrossentropy(),
                optimizer=opt,
                metrics=["accuracy"],
                steps_per_execution=8)

  model.fit(dataset, epochs=50, steps_per_epoch=8)