def main(argv=None):
    del argv
    hparams = create_hparams(FLAGS.experiment)
    if (hparams['distribution_strategy'] == 'multi_worker_mirrored'):
        strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
    elif (hparams['distribution_strategy'] == 'mirrored'):
        strategy = tf.distribute.MirroredStrategy()
    else:
        raise ValueError(('Only `multi_worker_mirrored` is supported strategy in Keras MNIST example at this time. Strategy passed in is %s' % hparams['distribution_strategy']))
    train_set = input_utils.input_fn(hparams['train_files'], hparams['batch_size'], hparams['vocab_size'], hparams['max_pixel_pos'], hparams['max_dom_pos'], epoches=1, buffer_size=hparams['train_buffer_size'])
    dev_set = input_utils.input_fn(hparams['eval_files'], hparams['eval_batch_size'], hparams['vocab_size'], hparams['max_pixel_pos'], hparams['max_dom_pos'], epoches=100, buffer_size=hparams['eval_buffer_size'])
    with strategy.scope():
        model = WidgetCaptionModel(hparams)
        lr_schedule = optimizer.LearningRateSchedule(hparams['learning_rate_constant'], hparams['hidden_size'], hparams['learning_rate_warmup_steps'])
        opt = tf.keras.optimizers.Adam(lr_schedule, hparams['optimizer_adam_beta1'], hparams['optimizer_adam_beta2'], epsilon=hparams['optimizer_adam_epsilon'])
        model.compile(optimizer=opt)
        init_resnet(hparams, model)
    callbacks = [tf.keras.callbacks.TerminateOnNaN()]
    if FLAGS.model_dir:
        ckpt_filepath = os.path.join(FLAGS.model_dir, 'saved/{epoch:04d}')
        backup_dir = os.path.join(FLAGS.model_dir, 'backup')
        tensorboard_callback = TensorBoardCallBack(log_dir=FLAGS.model_dir)
        callbacks.append(tensorboard_callback)
        model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_filepath, save_weights_only=True)
        callbacks.append(model_checkpoint_callback)
        if tf.executing_eagerly():
            callbacks.append(tf.keras.callbacks.experimental.BackupAndRestore(backup_dir=backup_dir))
    history = model.fit(x=train_set, epochs=hparams['train_epoches'], validation_data=dev_set, validation_steps=10, callbacks=callbacks)
    logging.info('Training ends successfully. `model.fit()` result: %s', history.history)
def main(argv=None):
    del argv
    hparams = create_hparams(FLAGS.experiment)
    if (hparams['distribution_strategy'] == 'multi_worker_mirrored'):
        strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
    elif (hparams['distribution_strategy'] == 'mirrored'):
        strategy = tf.distribute.MirroredStrategy()
    else:
        raise ValueError(('Only `multi_worker_mirrored` is supported strategy in Keras MNIST example at this time. Strategy passed in is %s' % hparams['distribution_strategy']))
    train_set = input_utils.input_fn(hparams['train_files'], hparams['batch_size'], hparams['vocab_size'], hparams['max_pixel_pos'], hparams['max_dom_pos'], epoches=1, buffer_size=hparams['train_buffer_size'])
    dev_set = input_utils.input_fn(hparams['eval_files'], hparams['eval_batch_size'], hparams['vocab_size'], hparams['max_pixel_pos'], hparams['max_dom_pos'], epoches=100, buffer_size=hparams['eval_buffer_size'])
    with strategy.scope():
        model = WidgetCaptionModel(hparams)
        lr_schedule = optimizer.LearningRateSchedule(hparams['learning_rate_constant'], hparams['hidden_size'], hparams['learning_rate_warmup_steps'])
        opt = tf.keras.optimizers.Adam(lr_schedule, hparams['optimizer_adam_beta1'], hparams['optimizer_adam_beta2'], epsilon=hparams['optimizer_adam_epsilon'])
        model.compile(optimizer=opt)
        init_resnet(hparams, model)
    callbacks = [tf.keras.callbacks.TerminateOnNaN()]
    if FLAGS.model_dir:
        ckpt_filepath = os.path.join(FLAGS.model_dir, 'saved/{epoch:04d}')
        backup_dir = os.path.join(FLAGS.model_dir, 'backup')
        tensorboard_callback = TensorBoardCallBack(log_dir=FLAGS.model_dir)
        callbacks.append(tensorboard_callback)
        model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_filepath, save_weights_only=True)
        callbacks.append(model_checkpoint_callback)
        if tf.executing_eagerly():
            callbacks.append(tf.keras.callbacks.experimental.BackupAndRestore(backup_dir=backup_dir))
    history = model.fit(x=train_set, epochs=hparams['train_epoches'], validation_data=dev_set, validation_steps=10, callbacks=callbacks)
    logging.info('Training ends successfully. `model.fit()` result: %s', history.history)