18. IPU TensorFlow Addons

18.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.

18.2. Keras layers

The ipu_tensorflow_addons.keras.layers namespace contains Keras layers optimised for running on IPUs. These layers can be used the same way as standard Keras layers.

Some IPU-specific versions of standard Keras layers are included. Swapping out standard keras layers for their IPU-specific counterparts will improve your model’s performance when using IPUs.

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

18.3. Optimizers

The optimizers contained in the IPU TensorFlow add-ons are drop in replacements to TensorFlow optimizers. They are functionally the same but have a number of additional features, which can be used via the optimizer’s kwargs.

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.