20. IPU TensorFlow Addons

20.1. Introduction

IPU TensorFlow Addons is a collection of addons created for IPU TensorFlow. These include layers and optimizers for Keras, as well as legacy TensorFlow layers and optimizers.

20.2. Keras layers

The ipu_tensorflow_addons.keras.layers namespace contains both IPU-specific implementations of standard Keras layers and IPU-specific layers that do not have upstream equivalents.

20.2.1. IPU implementations of standard Keras layers

Note

Swapping standard Keras layers for their IPU-specific counterparts will improve your model’s performance when using IPUs.

20.2.2. Layers without upstream equivalents

The layers below handle IPU specific behaviour:

The layers below are designed specifically for the IPU, but are not part of upstream Keras:

20.2.3. Code example

The example below demonstrates the use of the ipu_tensorflow_addons.keras.layers namespace in an IPU application.

 1import argparse
 2import tensorflow as tf
 3from tensorflow.python import ipu
 4from tensorflow import keras
 5
 6from ipu_tensorflow_addons.keras import layers as ipu_layers
 7
 8max_features = 20000
 9
10
11# Define the dataset
12def get_dataset():
13  (x_train,
14   y_train), (_, _) = keras.datasets.imdb.load_data(num_words=max_features)
15
16  x_train = keras.preprocessing.sequence.pad_sequences(x_train, maxlen=80)
17
18  ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
19  ds = ds.repeat()
20  ds = ds.map(lambda x, y: (x, tf.cast(y, tf.int32)))
21  ds = ds.batch(32, drop_remainder=True)
22  return ds
23
24
25# Define the model
26def get_model():
27  input_layer = keras.layers.Input(shape=(80), dtype=tf.int32, batch_size=32)
28
29  with keras.ipu.PipelineStage(0):
30    x = ipu_layers.Embedding(max_features, 64)(input_layer)
31    x = ipu_layers.LSTM(64, dropout=0.2)(x)
32
33  with keras.ipu.PipelineStage(1):
34    a = keras.layers.Dense(8, activation='relu')(x)
35
36  with keras.ipu.PipelineStage(2):
37    b = keras.layers.Dense(8, activation='relu')(x)
38
39  with keras.ipu.PipelineStage(3):
40    x = keras.layers.Concatenate()([a, b])
41    x = keras.layers.Dense(1, activation='sigmoid')(x)
42
43  return keras.Model(input_layer, x)
44
45
46#
47# Main code
48#
49
50# Parse command line args
51parser = argparse.ArgumentParser("Config Parser", add_help=False)
52parser.add_argument('--steps-per-epoch',
53                    type=int,
54                    default=768,
55                    help="Number of steps in each epoch.")
56parser.add_argument('--epochs',
57                    type=int,
58                    default=3,
59                    help="Number of epochs to run.")
60args = parser.parse_args()
61
62# Configure IPUs
63cfg = ipu.config.IPUConfig()
64cfg.auto_select_ipus = 2
65cfg.configure_ipu_system()
66
67# Set up IPU strategy
68strategy = ipu.ipu_strategy.IPUStrategyV1()
69with strategy.scope():
70
71  model = get_model()
72  model.set_pipelining_options(gradient_accumulation_steps_per_replica=8,
73                               device_mapping=[0, 1, 1, 0])
74  model.compile(loss='binary_crossentropy',
75                optimizer=keras.optimizers.Adam(0.005),
76                steps_per_execution=16)
77
78  model.fit(get_dataset(),
79            steps_per_epoch=args.steps_per_epoch,
80            epochs=args.epochs)

20.3. Optimizers

The ipu_tensorflow_addons.keras.optimizers namespace contains IPU-specific implementations of Keras optimizers from TensorFlow Addons and upstream Keras:

They are functionally the same but have a number of additional features, which can be used via the optimizer’s keyword arguments.

The precision of any optimizer states within the optimizer can be set independently of each other and the model parameters. This is particularly useful when training in mixed precision.

The optimizer update can be outlined, making the optimizer update block code reusable, which can reduce memory at the expense of passing variables around.