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 both IPU-specific implementations
of standard Keras layers and IPU-specific layers that do not have upstream equivalents.
18.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.
18.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:
18.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
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 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.