def main(argv):
del argv # unused
logging.log(logging.INFO, "Logging application {}".format(__file__))
if FLAGS.debug:
logging.set_verbosity(logging.DEBUG)
logging.log(logging.DEBUG, "Running in debug mode")
physical_devices = tf.config.experimental.list_physical_devices("GPU")
tf.config.experimental.set_memory_growth(physical_devices[0], True)
# create output directory if none specified
if FLAGS.output_dir is None:
output_dir = os.path.join(
"logs", __file__,
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
file_io.check_create_dir(output_dir)
# output directory specified, load model options if found
else:
output_dir = FLAGS.output_dir
# print flag options
flag_options = {}
for flag in FLAGS.get_key_flags_for_module(__file__):
flag_options[flag.name] = flag.value
# logging
logging_utils.absl_file_logger(output_dir, f"log.test")
logging.log(logging.INFO, f"Model directory: {output_dir}")
logging.log(logging.INFO, f"Flag options: {flag_options}")
# set seeds for reproducibility
np.random.seed(FLAGS.seed)
tf.random.set_seed(FLAGS.seed)
# test baseline matching model (no background training step)
test()
def main(argv):
"""Main program logic."""
del argv # unused
logging.log(logging.INFO, "Logging application {}".format(__file__))
if FLAGS.debug:
logging.set_verbosity(logging.DEBUG)
logging.log(logging.DEBUG, "Running in debug mode")
physical_devices = tf.config.experimental.list_physical_devices("GPU")
tf.config.experimental.set_memory_growth(physical_devices[0], True)
model_found = False
# no prior run specified, train model
if FLAGS.output_dir is None:
if FLAGS.target != "train":
raise ValueError(
f"Target `{FLAGS.target}` requires --output_dir to be specified.")
output_dir = os.path.join(
"logs", __file__, datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
file_io.check_create_dir(output_dir)
model_options = DEFAULT_OPTIONS
# add flag options to model options
model_options["base_dir"] = FLAGS.base_dir
if FLAGS.base_dir is None:
raise ValueError(
f"Target `{FLAGS.target}` requires --base_dir to be specified.")
model_options["base_model"] = (
"best_model" if FLAGS.load_best else "model")
# prior run specified, resume training or test model
else:
output_dir = FLAGS.output_dir
# load current or best model
model_file = "best_model.h5" if FLAGS.load_best else "model.h5"
model_step_file = "best_model.step" if FLAGS.load_best else "model.step"
if FLAGS.base_dir is not None:
raise ValueError(
f"Flag --base_dir should not be set for target `{FLAGS.target}`.")
if os.path.exists(os.path.join(output_dir, model_file)):
model_found = True
model_options = file_io.read_json(
os.path.join(output_dir, "model_options.json"))
elif FLAGS.target != "train":
raise ValueError(
f"Target `{FLAGS.target}` specified but `{model_file}` not "
f"found in {output_dir}.")
# gather flag options
flag_options = {}
for flag in FLAGS.get_key_flags_for_module(__file__):
flag_options[flag.name] = flag.value
# logging
logging_utils.absl_file_logger(output_dir, f"log.{FLAGS.target}")
logging.log(logging.INFO, f"Model directory: {output_dir}")
logging.log(logging.INFO, f"Model options: {model_options}")
logging.log(logging.INFO, f"Flag options: {flag_options}")
tf_writer = None
if FLAGS.tensorboard and FLAGS.target == "train":
tf_writer = tf.summary.create_file_writer(output_dir)
# set seeds for reproducibility
np.random.seed(model_options["seed"])
tf.random.set_seed(model_options["seed"])
# run target
if FLAGS.target == "train":
if model_found and FLAGS.resume:
train(model_options, output_dir, model_file, model_step_file,
tf_writer=tf_writer)
else:
train(model_options, output_dir, tf_writer=tf_writer)
elif FLAGS.target == "validate": # TODO
raise NotImplementedError
elif FLAGS.target == "embed":
embed(model_options, output_dir, model_file, model_step_file)
else:
test(model_options, output_dir, model_file, model_step_file)