def run_model(opts):
# Use Keras to get the dataset:
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
# Sizes/shapes for the dataset:
image_shape = x_train.shape[1:]
num_pixels = image_shape[0] * image_shape[1]
batch_size = 16
batch_shape = [batch_size, num_pixels]
num_train = y_train.shape[0]
num_test = y_test.shape[0]
data_shape = [None, num_pixels]
# Flatten the images and cast the labels:
x_train_flat = x_train.astype(np.float32).reshape(-1, num_pixels)
x_test_flat = x_test.astype(np.float32).reshape(-1, num_pixels)
y_train = y_train.astype(np.int32)
y_test = y_test.astype(np.int32)
# Decide how to split epochs into loops up front:
epochs = 5
ipu_steps_per_epoch = 15
batches_per_epoch = num_train // batch_size
train_batches = (num_train * epochs) // batch_size
test_batches = num_test // batch_size
batches_per_step = batches_per_epoch // ipu_steps_per_epoch
if not batches_per_epoch % ipu_steps_per_epoch == 0:
raise ValueError(f"IPU steps per epoch {ipu_steps_per_epoch} must divide batches per epoch {batches_per_epoch}.")
# Put placeholders on the CPU host:
with tf.device("cpu"):
place_x = tf.placeholder(dtype=tf.float32, shape=data_shape, name="input")
place_y = tf.placeholder(dtype=tf.int32, shape=[None], name="label")
lr_placeholder = tf.placeholder(tf.float32, shape=[])
# Create dataset and IPU feeds:
dataset = tf.data.Dataset.from_tensor_slices((place_x, place_y))
dataset = dataset.cache().repeat().batch(batch_size, drop_remainder=True)
infeed_train_queue = ipu_infeed_queue.IPUInfeedQueue(
dataset, feed_name="train_infeed")
outfeed_train_queue = ipu_outfeed_queue.IPUOutfeedQueue(
feed_name="train_outfeed")
infeed_test_queue = ipu_infeed_queue.IPUInfeedQueue(
dataset, feed_name="test_infeed")
outfeed_test_queue = ipu_outfeed_queue.IPUOutfeedQueue(
feed_name="test_outfeed")
# Use function binding to create all the builder functions that are neeeded:
bound_train_model = partial(model, lr_placeholder, outfeed_train_queue, True)
bound_train_loop = partial(
loop_builder, batches_per_step, bound_train_model, infeed_train_queue)
bound_test_model = partial(model, lr_placeholder, outfeed_test_queue, False)
bound_test_loop = partial(loop_builder, test_batches,
bound_test_model, infeed_test_queue)
# Use the bound builder functions to place the model on the IPU:
with scopes.ipu_scope("/device:IPU:0"):
train_loop = ipu_compiler.compile(bound_train_loop, inputs=[])
test_loop = ipu_compiler.compile(bound_test_loop, inputs=[])
# Initialisers should go on the CPU:
with tf.device("cpu"):
metrics_vars = tf.get_collection(
tf.GraphKeys.LOCAL_VARIABLES, scope="metrics")
metrics_initializer = tf.variables_initializer(var_list=metrics_vars)
saver = tf.train.Saver()
# Setup and acquire an IPU device:
config = utils.create_ipu_config()
config = utils.auto_select_ipus(config, 1)
utils.configure_ipu_system(config)
# These allow us to retrieve the results of IPU feeds:
dequeue_train_outfeed = outfeed_train_queue.dequeue()
dequeue_test_outfeed = outfeed_test_queue.dequeue()
# Create a benchmark program for the infeed to determine maximum achievable throughput:
infeed_perf = dataset_benchmark.infeed_benchmark(
infeed_train_queue, epochs, num_train, True)
print(f"\nImage shape: {image_shape} Training examples: {num_train} Test examples: {num_test}")
print(f"Epochs: {epochs} Batch-size: {batch_size} Steps-per-epoch: {ipu_steps_per_epoch} Batches-per-step: {batches_per_step}")
# Run the model:
with tf.Session() as sess:
print(f"Benchmarking the infeed...")
sess.run(infeed_perf, feed_dict={place_x: x_train_flat, place_y: y_train})
sess.run(tf.global_variables_initializer())
sess.run(infeed_train_queue.initializer, feed_dict={
place_x: x_train_flat, place_y: y_train})
if opts.test_mode in ["all", "training"]:
print(f"Training...")
progress = tqdm(
range(epochs), bar_format='{desc} Epoch: {n_fmt}/{total_fmt} {bar}')
for e in progress:
sess.run(metrics_initializer)
for i in range(ipu_steps_per_epoch):
sess.run(train_loop, feed_dict={lr_placeholder: scheduler(e)})
result = sess.run(dequeue_train_outfeed)
if len(result['mean_loss'] != 0) and len(result['acc'] != 0):
progress.set_description(f"Loss {result['mean_loss'][0]:.5f} Accuracy {result['acc'][0]:.5f}")
print(f"Saving...")
saver.save(sess, "model")
if opts.test_mode in ["all", "tests"]:
print(f"Testing...")
sess.run(metrics_initializer)
sess.run(infeed_test_queue.initializer, feed_dict={
place_x: x_test_flat, place_y: y_test})
sess.run(test_loop)
result = sess.run(dequeue_test_outfeed)
test_loss = np.mean(result['mean_loss'])
test_acc = np.mean(result['acc'])
print(f"Test loss: {test_loss:.8f} Test accuracy: {test_acc:.8f}")
metrics_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES,
scope="metrics")
metrics_initializer = tf.variables_initializer(var_list=metrics_vars)
saver = tf.train.Saver()
# Setup and acquire an IPU device:
config = utils.create_ipu_config()
config = utils.auto_select_ipus(config, 1)
utils.configure_ipu_system(config)
# These allow us to retrieve the results of IPU feeds:
dequeue_train_outfeed = outfeed_train_queue.dequeue()
dequeue_test_outfeed = outfeed_test_queue.dequeue()
# Create a benchmark program for the infeed to determine maximum achievable throughput:
infeed_perf = dataset_benchmark.infeed_benchmark(infeed_train_queue, epochs,
num_train, True)
print(
f"\nImage shape: {image_shape} Training examples: {num_train} Test examples: {num_test}"
)
print(
f"Epochs: {epochs} Batch-size: {batch_size} Steps-per-epoch: {ipu_steps_per_epoch} Batches-per-step: {batches_per_step}"
)
# Run the model:
with tf.Session() as sess:
print(f"Benchmarking the infeed...")
sess.run(infeed_perf, feed_dict={place_x: x_train_flat, place_y: y_train})
sess.run(tf.global_variables_initializer())
sess.run(infeed_train_queue.initializer,
datefmt='%Y-%m-%d %H:%M:%S')
# Parse options
opts = parse_args()
if not opts.on_device_only:
logger.info("Creating training dataset, infeed queue and benchmark.")
# Create training dataset and infeed queue
train_set, num_train, *_ = make_dataset(opts, use_synthetic_data=False, training=True)
num_train = num_train // opts.batch_size
infeed_train_queue = ipu_infeed_queue.IPUInfeedQueue(train_set)
# Benchmark it
infeed_perf_train = dataset_benchmark.infeed_benchmark(
infeed_queue=infeed_train_queue,
number_of_epochs=opts.epochs,
elements_per_epochs=num_train,
print_stats=False)
ds_perf_train = dataset_benchmark.dataset_benchmark(
dataset=train_set,
number_of_epochs=opts.epochs,
elements_per_epochs=num_train,
print_stats=False,
apply_options=True)
logger.info("Creating test dataset, infeed queue and benchmark.")
# Create test dataset
test_set, num_test, *_ = make_dataset(opts, use_synthetic_data=False, training=False)
num_test = num_test // opts.batch_size
infeed_test_queue = ipu_infeed_queue.IPUInfeedQueue(test_set)
def train(*tf_records: "Records to train on"):
"""Train on examples."""
tf.logging.set_verbosity(tf.logging.INFO)
estimator = dual_net.get_estimator()
effective_batch_size = FLAGS.train_batch_size
if FLAGS.use_tpu:
effective_batch_size *= FLAGS.num_tpu_cores
elif FLAGS.use_ipu:
effective_batch_size *= FLAGS.num_ipu_cores
if FLAGS.use_tpu:
if FLAGS.use_bt:
def _input_fn(params):
games = bigtable_input.GameQueue(FLAGS.cbt_project,
FLAGS.cbt_instance,
FLAGS.cbt_table)
games_nr = bigtable_input.GameQueue(FLAGS.cbt_project,
FLAGS.cbt_instance,
FLAGS.cbt_table + '-nr')
return preprocessing.get_tpu_bt_input_tensors(
games,
games_nr,
params['batch_size'],
number_of_games=FLAGS.window_size,
random_rotation=True)
else:
def _input_fn(params):
return preprocessing.get_tpu_input_tensors(
params['batch_size'], tf_records, random_rotation=True)
# Hooks are broken with TPUestimator at the moment.
hooks = []
elif FLAGS.use_ipu:
def _input_fn():
return preprocessing.get_ipu_input_tensors(
FLAGS.train_batch_size,
tf_records,
filter_amount=FLAGS.filter_amount,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
shuffle_examples=False,
random_rotation=False)
hooks = []
else:
def _input_fn():
return preprocessing.get_input_tensors(
FLAGS.train_batch_size,
tf_records,
filter_amount=FLAGS.filter_amount,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
random_rotation=True)
hooks = [
UpdateRatioSessionHook(FLAGS.work_dir),
EchoStepCounterHook(output_dir=FLAGS.work_dir)
]
try:
if FLAGS.PROFILING:
ph = ProfilerHook()
hooks = [ph]
except:
pass
steps = FLAGS.steps_to_train
# step correction due to smaller batch size
if FLAGS.use_ipu:
steps = steps * 4096 // effective_batch_size
logging.info("Training, steps = %s, batch = %s -> %s examples", steps
or '?', effective_batch_size,
(steps * effective_batch_size) if steps else '?')
if FLAGS.use_bt:
games = bigtable_input.GameQueue(FLAGS.cbt_project, FLAGS.cbt_instance,
FLAGS.cbt_table)
if not games.read_wait_cell():
games.require_fresh_games(20000)
latest_game = games.latest_game_number
index_from = max(latest_game, games.read_wait_cell())
print("== Last game before training:", latest_game, flush=True)
print("== Wait cell:", games.read_wait_cell(), flush=True)
if DATA_BENCHMARK:
benchmark_op = dataset_benchmark(
dataset=_input_fn(),
number_of_epochs=80,
elements_per_epochs=10000,
print_stats=True,
# apply_options=False
)
import json
print("Benchmarking data pipeline:")
with tf.Session() as sess:
json_string = sess.run(benchmark_op)
json_object = json.loads(json_string[0])
print(json_object)
if not INFEED_BENCHMARK:
raise NotImplementedError("Data benchmark ended.")
else:
print("Data benchmark ended.")
if INFEED_BENCHMARK:
benchmark_op = infeed_benchmark(
infeed_queue=ipu_infeed_queue.IPUInfeedQueue(_input_fn(),
feed_name="infeed"),
number_of_epochs=80,
elements_per_epochs=10000,
print_stats=True,
# apply_options=False
)
import json
print("Benchmarking data pipeline:")
with tf.Session() as sess:
json_string = sess.run(benchmark_op)
json_object = json.loads(json_string[0])
print(json_object)
raise NotImplementedError("Infeed benchmark ended.")
try:
estimator.train(_input_fn, steps=steps, hooks=hooks)
if FLAGS.use_bt:
bigtable_input.set_fresh_watermark(games, index_from,
FLAGS.window_size)
except:
if FLAGS.use_bt:
games.require_fresh_games(0)
raise
return estimator