Eager execution workflow with RevNet trained on CIFAR-10.
(_)
| 35 | |
| 36 | |
| 37 | def main(_): |
| 38 | """Eager execution workflow with RevNet trained on CIFAR-10.""" |
| 39 | tf.enable_eager_execution() |
| 40 | |
| 41 | config = get_config(config_name=FLAGS.config, dataset=FLAGS.dataset) |
| 42 | ds_train, ds_train_one_shot, ds_validation, ds_test = get_datasets( |
| 43 | data_dir=FLAGS.data_dir, config=config) |
| 44 | model = revnet.RevNet(config=config) |
| 45 | global_step = tf.train.get_or_create_global_step() # Ensure correct summary |
| 46 | global_step.assign(1) |
| 47 | learning_rate = tf.train.piecewise_constant( |
| 48 | global_step, config.lr_decay_steps, config.lr_list) |
| 49 | optimizer = tf.train.MomentumOptimizer( |
| 50 | learning_rate, momentum=config.momentum) |
| 51 | checkpointer = tf.train.Checkpoint( |
| 52 | optimizer=optimizer, model=model, optimizer_step=global_step) |
| 53 | |
| 54 | if FLAGS.use_defun: |
| 55 | model.call = tfe.defun(model.call) |
| 56 | model.compute_gradients = tfe.defun(model.compute_gradients) |
| 57 | model.get_moving_stats = tfe.defun(model.get_moving_stats) |
| 58 | model.restore_moving_stats = tfe.defun(model.restore_moving_stats) |
| 59 | global apply_gradients # pylint:disable=global-variable-undefined |
| 60 | apply_gradients = tfe.defun(apply_gradients) |
| 61 | |
| 62 | if FLAGS.train_dir: |
| 63 | summary_writer = tf.contrib.summary.create_file_writer(FLAGS.train_dir) |
| 64 | if FLAGS.restore: |
| 65 | latest_path = tf.train.latest_checkpoint(FLAGS.train_dir) |
| 66 | checkpointer.restore(latest_path) |
| 67 | print("Restored latest checkpoint at path:\"{}\" " |
| 68 | "with global_step: {}".format(latest_path, global_step.numpy())) |
| 69 | sys.stdout.flush() |
| 70 | |
| 71 | for x, y in ds_train: |
| 72 | train_one_iter(model, x, y, optimizer, global_step=global_step) |
| 73 | |
| 74 | if global_step.numpy() % config.log_every == 0: |
| 75 | acc_test, loss_test = evaluate(model, ds_test) |
| 76 | |
| 77 | if FLAGS.validate: |
| 78 | acc_train, loss_train = evaluate(model, ds_train_one_shot) |
| 79 | acc_validation, loss_validation = evaluate(model, ds_validation) |
| 80 | print("Iter {}, " |
| 81 | "training set accuracy {:.4f}, loss {:.4f}; " |
| 82 | "validation set accuracy {:.4f}, loss {:.4f}; " |
| 83 | "test accuracy {:.4f}, loss {:.4f}".format( |
| 84 | global_step.numpy(), acc_train, loss_train, acc_validation, |
| 85 | loss_validation, acc_test, loss_test)) |
| 86 | else: |
| 87 | print("Iter {}, test accuracy {:.4f}, loss {:.4f}".format( |
| 88 | global_step.numpy(), acc_test, loss_test)) |
| 89 | sys.stdout.flush() |
| 90 | |
| 91 | if FLAGS.train_dir: |
| 92 | with summary_writer.as_default(): |
| 93 | with tf.contrib.summary.always_record_summaries(): |
| 94 | tf.contrib.summary.scalar("Test accuracy", acc_test) |
nothing calls this directly
no test coverage detected