5. TensorFlow 2 examples

5.1. Training on the IPU

This example shows how to use the IPU-specific Keras Model class and the IPUStrategy to train a model using the Keras Model.fit() method.

The IPU specific changes are highlighted:

  • Import the IPU extensions to TensorFlow.

  • Create a configuration for the IPU target. To keep things simple, this just selects the first available IPU. A configuration can select specific IPUs, or groups of IPUs.

  • Call the IPU implementation of Sequential(). This has exactly the same interface as the standard Keras implementation but uses layers that are optimised for the IPU.

  • Use the IPU distribution strategy as a context in order to place the training code on the configured IPU.

As you can see, there are minimal changes required to get a Keras model running on the IPU. No changes at all are required to the layers, other than using the IPU-optimised versions.

 1from __future__ import absolute_import, division, print_function, unicode_literals
 2
 3import tensorflow as tf
 4
 5from tensorflow import keras
 6from tensorflow.python import ipu
 7
 8#
 9# Configure the IPU system
10#
11cfg = ipu.utils.create_ipu_config()
12cfg = ipu.utils.auto_select_ipus(cfg, 1)
13ipu.utils.configure_ipu_system(cfg)
14
15
16#
17# Create the input data and labels
18#
19def create_dataset():
20  mnist = tf.keras.datasets.mnist
21
22  (x_train, y_train), (_, _) = mnist.load_data()
23  x_train = x_train / 255.0
24
25  train_ds = tf.data.Dataset.from_tensor_slices(
26      (x_train, y_train)).shuffle(10000).batch(32, drop_remainder=True)
27  train_ds = train_ds.map(lambda d, l:
28                          (tf.cast(d, tf.float32), tf.cast(l, tf.float32)))
29
30  return train_ds.repeat()
31
32
33#
34# Create the model using the IPU-specific Sequential class
35#
36def create_model():
37  m = ipu.keras.Sequential([
38      keras.layers.Flatten(),
39      keras.layers.Dense(128, activation='relu'),
40      keras.layers.Dense(10, activation='softmax')
41  ])
42  return m
43
44
45# Create an IPU distribution strategy
46strategy = ipu.ipu_strategy.IPUStrategy()
47
48with strategy.scope():
49  # Create an instance of the model
50  model = create_model()
51
52  # Get the training dataset
53  ds = create_dataset()
54
55  # Train the model
56  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
57                optimizer=tf.keras.optimizers.SGD())
58  model.fit(ds, steps_per_epoch=2000, epochs=4)

5.2. Custom training function

This example shows how to use a custom training function with the IPUStrategy and the standard Keras Sequential class.

 1from __future__ import absolute_import, division, print_function, unicode_literals
 2
 3import tensorflow as tf
 4
 5from tensorflow import keras
 6from tensorflow.python import ipu
 7
 8step_count = 10000
 9
10#
11# Configure the IPU system
12#
13cfg = ipu.utils.create_ipu_config()
14cfg = ipu.utils.auto_select_ipus(cfg, 1)
15ipu.utils.configure_ipu_system(cfg)
16
17
18#
19# The input data and labels
20#
21def create_dataset():
22  mnist = tf.keras.datasets.mnist
23
24  (x_train, y_train), (_, _) = mnist.load_data()
25  x_train = x_train / 255.0
26
27  train_ds = tf.data.Dataset.from_tensor_slices(
28      (x_train, y_train)).shuffle(10000).batch(32)
29  train_ds = train_ds.map(lambda d, l:
30                          (tf.cast(d, tf.float32), tf.cast(l, tf.int32)))
31
32  return train_ds.repeat()
33
34
35#
36# The model.  Because this model does not have a specific shape for its inputs
37# it will be constructed when it is first called (in the `train` function). So
38# it does not need to be an IPU device targeted model.
39#
40def create_model():
41  m = keras.models.Sequential([
42      keras.layers.Flatten(),
43      keras.layers.Dense(128, activation='relu'),
44      keras.layers.Dense(10, activation='softmax')
45  ])
46  return m
47
48
49# The custom training loop
50@tf.function
51def training_step(features, labels, model, opt):
52  with tf.GradientTape() as tape:
53    predictions = model(features, training=True)
54    prediction_loss = keras.losses.sparse_categorical_crossentropy(
55        labels, predictions)
56    loss = tf.reduce_mean(prediction_loss)
57
58  grads = tape.gradient(loss, model.trainable_variables)
59  opt.apply_gradients(zip(grads, model.trainable_variables))
60  return loss
61
62
63# Create an IPU distribution strategy
64strategy = ipu.ipu_strategy.IPUStrategy()
65
66with strategy.scope():
67  # An optimizer for updating the trainable variables
68  opt = tf.keras.optimizers.SGD(0.01)
69
70  # Create an instance of the model
71  model = create_model()
72
73  # Get the training dataset
74  ds = create_dataset()
75
76  # Train the model
77  for (x, y), c in zip(ds, range(step_count)):
78    loss = strategy.experimental_run_v2(training_step, args=[x, y, model, opt])
79
80    if not c % 50:
81      print("Step " + str(c) + " loss = " + str(loss.numpy()))

5.3. Pipelined model

This example shows how to use the IPU-specific Keras pipelined Model class to train a network.

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

This example shows how to use the IPU-specific Keras pipelined Sequential class to train a network.

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