def train_fn(hparams, num_workers): """Copy of train function from estimator.py.""" # TODO: Merge improvements into the original. # pylint: disable=protected-access hparams.tgt_sos_id, hparams.tgt_eos_id = nmt_estimator._get_tgt_sos_eos_id( hparams) model_fn = nmt_estimator.make_model_fn(hparams) def print_log(): mlperf_log.gnmt_print(key=mlperf_log.TRAIN_LOOP) mlperf_log.gnmt_print(key=mlperf_log.TRAIN_EPOCH, value=0) mlperf_log.gnmt_print(key=mlperf_log.INPUT_SIZE, value=hparams.num_examples_per_epoch) if hparams.use_tpu_low_level_api: runner = create_train_runner(hparams, num_workers) mlperf_log.gnmt_print(key=mlperf_log.RUN_START) input_fn = DistributedPipeline(hparams, num_workers) runner.initialize(input_fn, {}) runner.build_model(model_fn, {}) print_log() runner.train(0, hparams.num_train_steps) return 0.0 # cluster = tf.contrib.cluster_resolver.TPUClusterResolver(hparams.tpu_name) # cluster_spec = cluster.cluster_spec() # print('cluster_spec: %s' % cluster_spec) # num_workers = cluster_spec.num_tasks('tpu_worker') # print('num_workers: %s' % num_workers) pipeline = DistributedPipeline(hparams, num_workers) print_log() if hparams.use_tpu: run_config = nmt_estimator._get_tpu_run_config(hparams, True) estimator = tf.contrib.tpu.TPUEstimator( model_fn=model_fn, config=run_config, use_tpu=hparams.use_tpu, train_batch_size=hparams.batch_size, eval_batch_size=hparams.batch_size, predict_batch_size=hparams.infer_batch_size) else: raise ValueError("Distributed input pipeline only supported on TPUs.") hooks = [pipeline] if hparams.use_async_checkpoint: hooks.append( async_checkpoint.AsyncCheckpointSaverHook( checkpoint_dir=hparams.out_dir, save_steps=int(hparams.num_examples_per_epoch / hparams.batch_size))) estimator.train(input_fn=pipeline, max_steps=hparams.num_train_steps, hooks=hooks) # Return value is not used return 0.0
def train_and_eval_fn(hparams, num_workers): """Train and evaluation function.""" # pylint: disable=protected-access mlperf_log.gnmt_print(key=mlperf_log.RUN_START) hparams.tgt_sos_id, hparams.tgt_eos_id = 1, 2 model_fn = nmt_estimator.make_model_fn(hparams) pipeline = DistributedPipeline(hparams, num_workers) run_config = nmt_estimator._get_tpu_run_config(hparams) estimator = tf.contrib.tpu.TPUEstimator( model_fn=model_fn, config=run_config, use_tpu=hparams.use_tpu, train_batch_size=hparams.batch_size, eval_batch_size=hparams.batch_size, predict_batch_size=hparams.infer_batch_size) score = 0.0 mlperf_log.gnmt_print(key=mlperf_log.TRAIN_LOOP) mlperf_log.gnmt_print(key=mlperf_log.EVAL_TARGET, value=hparams.target_bleu) for i in range(hparams.max_train_epochs): mlperf_log.gnmt_print(key=mlperf_log.TRAIN_EPOCH, value=i) tf.logging.info("Start training epoch %d", i) mlperf_log.gnmt_print(key=mlperf_log.INPUT_SIZE, value=hparams.num_examples_per_epoch) steps_per_epoch = int(hparams.num_examples_per_epoch / hparams.batch_size) max_steps = steps_per_epoch * (i + 1) estimator.train(input_fn=pipeline, max_steps=max_steps, hooks=[pipeline]) mlperf_log.gnmt_print(key=mlperf_log.TRAIN_CHECKPOINT, value=("Under " + hparams.out_dir)) tf.logging.info("End training epoch %d", i) mlperf_log.gnmt_print(key=mlperf_log.EVAL_START) score = nmt_estimator.get_metric_from_estimator(hparams, estimator) tf.logging.info("Score after epoch %d: %f", i, score) mlperf_log.gnmt_print(key=mlperf_log.EVAL_ACCURACY, value={ "value": score, "epoch": i }) mlperf_log.gnmt_print(key=mlperf_log.EVAL_STOP, value=i) if score >= hparams.target_bleu: mlperf_log.gnmt_print(mlperf_log.RUN_STOP, {"success": True}) return score mlperf_log.gnmt_print(mlperf_log.RUN_STOP, {"success": False}) return score