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