TensorFlow 2 examples

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.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
from tensorflow.python.ipu.config import IPUConfig

import tensorflow as tf

from tensorflow import keras
from tensorflow.python import ipu

#
# Configure the IPU system
#
cfg = IPUConfig()
cfg.auto_select_ipus = 1
cfg.configure_ipu_system()


#
# Create the input data and labels
#
def create_dataset():
  mnist = tf.keras.datasets.mnist

  (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.float32)))

  return train_ds.repeat()


#
# Create the model using the IPU-specific Sequential class
#
def create_model():
  m = keras.Sequential([
      keras.layers.Flatten(),
      keras.layers.Dense(128, activation='relu'),
      keras.layers.Dense(10, activation='softmax')
  ])
  return m


# Create an IPU distribution strategy
strategy = ipu.ipu_strategy.IPUStrategyV1()

with strategy.scope():
  # Create an instance of the model
  model = create_model()

  # Get the training dataset
  ds = create_dataset()

  # Train the model
  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                optimizer=tf.keras.optimizers.SGD(),
                steps_per_execution=100)
  model.fit(ds, steps_per_epoch=2000, epochs=4)

Custom training function

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

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf

from tensorflow import keras
from tensorflow.python import ipu

step_count = 1000
steps_per_execution = 10

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


#
# The input data and labels.
#
def create_dataset():
  mnist = tf.keras.datasets.mnist

  (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()


#
# The model. Because this model does not have a specific shape for its inputs
# it will be constructed when it is first called (in the `training_step`
# function).
#
def create_model():
  m = keras.models.Sequential([
      keras.layers.Flatten(),
      keras.layers.Dense(128, activation='relu'),
      keras.layers.Dense(10, activation='softmax')
  ])
  return m


#
# The custom training loop.
#
@tf.function(experimental_compile=True)
def training_loop(iterator, outfeed_queue, model, optimizer, num_iterations):
  for _ in tf.range(num_iterations):
    # Get the data for the step.
    features, labels = next(iterator)

    # Perform the training step.
    with tf.GradientTape() as tape:
      predictions = model(features, training=True)
      prediction_loss = keras.losses.sparse_categorical_crossentropy(
          labels, predictions)
      loss = tf.reduce_mean(prediction_loss)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    # Store the loss in the outfeed queue.
    outfeed_queue.enqueue(loss)


# Create an IPU distribution strategy
strategy = ipu.ipu_strategy.IPUStrategyV1()

with strategy.scope():
  # An optimizer for updating the trainable variables.
  opt = tf.keras.optimizers.SGD(0.01)

  # Create an instance of the model.
  model = create_model()

  # Create an iterator for the dataset.
  iterator = iter(create_dataset())

  # Create an IPUOutfeedQueue to collect results from each on device step.
  outfeed_queue = ipu.ipu_outfeed_queue.IPUOutfeedQueue()

  # Train the model.
  for step_begin in range(0, step_count, steps_per_execution):

    # Run `steps_per_execution` at a time.
    strategy.run(
        training_loop,
        args=[iterator, outfeed_queue, model, opt, steps_per_execution])

    # Get results for each step.
    for step, loss in zip(range(step_begin, step_begin + steps_per_execution),
                          outfeed_queue):
      if step % 50 == 0:
        print(f"Step {step}: loss {loss}")

Pipelined model

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import argparse
import tensorflow as tf

from tensorflow.python import ipu
from tensorflow.python.ipu.keras import layers as ipu_layers

from tensorflow.python.keras.datasets import imdb
from tensorflow.python import keras
from tensorflow.python.keras import layers
from tensorflow.python.keras.preprocessing import sequence
from tensorflow.python.keras.optimizer_v2.adam import Adam

max_features = 20000


# Define the dataset
def get_dataset():
  (x_train, y_train), (_, _) = imdb.load_data(num_words=max_features)

  x_train = sequence.pad_sequences(x_train, maxlen=80)

  ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
  ds = ds.repeat()
  ds = ds.map(lambda x, y: (x, tf.cast(y, tf.int32)))
  ds = ds.batch(32, drop_remainder=True)
  return ds


# Define the model
def get_model():
  input_layer = layers.Input(shape=(80), dtype=tf.int32, batch_size=32)

  with ipu.keras.PipelineStage(0):
    x = ipu_layers.Embedding(max_features, 64)(input_layer)
    x = ipu_layers.LSTM(64, dropout=0.2)(x)

  with ipu.keras.PipelineStage(1):
    a = layers.Dense(8, activation='relu')(x)

  with ipu.keras.PipelineStage(2):
    b = layers.Dense(8, activation='relu')(x)

  with ipu.keras.PipelineStage(3):
    x = layers.Concatenate()([a, b])
    x = layers.Dense(1, activation='sigmoid')(x)

  return keras.Model(input_layer, x)


#
# Main code
#

# Parse command line args
parser = argparse.ArgumentParser("Config Parser", add_help=False)
parser.add_argument('--steps-per-epoch',
                    type=int,
                    default=768,
                    help="Number of steps in each epoch.")
parser.add_argument('--epochs',
                    type=int,
                    default=3,
                    help="Number of epochs to run.")
args = parser.parse_args()

# Configure IPUs
cfg = ipu.config.IPUConfig()
cfg.auto_select_ipus = 2
cfg.configure_ipu_system()

# Set up IPU strategy
strategy = ipu.ipu_strategy.IPUStrategyV1()
with strategy.scope():

  model = get_model()
  model.set_pipelining_options(gradient_accumulation_steps_per_replica=8,
                               device_mapping=[0, 1, 1, 0])
  model.compile(loss='binary_crossentropy',
                optimizer=Adam(0.005),
                steps_per_execution=16)

  model.fit(get_dataset(),
            steps_per_epoch=args.steps_per_epoch,
            epochs=args.epochs)

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import argparse
import tensorflow as tf

from tensorflow.python import ipu

from tensorflow.python.ipu.keras.layers import Embedding
from tensorflow.python.ipu.keras.layers import LSTM

from tensorflow.python.keras.datasets import imdb
from tensorflow.python.keras.layers import Dense
from tensorflow.python.keras.preprocessing import sequence
from tensorflow.python.keras.optimizer_v2.adam import Adam

max_features = 20000


# Define the dataset
def get_dataset():
  (x_train, y_train), (_, _) = imdb.load_data(num_words=max_features)

  x_train = sequence.pad_sequences(x_train, maxlen=80)

  ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
  ds = ds.repeat()
  ds = ds.map(lambda x, y: (x, tf.cast(y, tf.int32)))
  ds = ds.batch(32, drop_remainder=True)
  return ds


# Define the model
def get_model():
  return ipu.keras.PipelineSequential(
      [[Embedding(max_features, 128)],
       [LSTM(128, dropout=0.2),
        Dense(1, activation='sigmoid')]],
      gradient_accumulation_count=16)


#
# Main code
#

# Parse command line args
parser = argparse.ArgumentParser("Config Parser", add_help=False)
parser.add_argument('--steps-per-epoch',
                    type=int,
                    default=768,
                    help="Number of steps in each epoch.")
parser.add_argument('--epochs',
                    type=int,
                    default=10,
                    help="Number of epochs to run.")
args = parser.parse_args()

# Configure IPUs
cfg = ipu.config.IPUConfig()
cfg.auto_select_ipus = 2
cfg.configure_ipu_system()

# Set up IPU strategy
strategy = ipu.ipu_strategy.IPUStrategyV1()
with strategy.scope():

  model = get_model()

  model.compile(loss='binary_crossentropy', optimizer=Adam(0.005))

  model.fit(get_dataset(),
            steps_per_epoch=args.steps_per_epoch,
            epochs=args.epochs)