def _run_train(ppo_hparams,
               event_dir,
               model_dir,
               restarter,
               train_summary_op,
               eval_summary_op,
               initializers,
               report_fn=None,
               model_save_fn=None):
  """Train."""
  summary_writer = tf.summary.FileWriter(
      event_dir, graph=tf.get_default_graph(), flush_secs=60)

  model_saver = tf.train.Saver(
      tf.global_variables(ppo_hparams.policy_network + "/.*") +
      tf.global_variables("training/" + ppo_hparams.policy_network + "/.*") +
      # tf.global_variables("clean_scope.*") +  # Needed for sharing params.
      tf.global_variables("global_step") +
      tf.global_variables("losses_avg.*") +
      tf.global_variables("train_stats.*")
  )

  global_step = tf.train.get_or_create_global_step()
  with tf.control_dependencies([tf.assign_add(global_step, 1)]):
    train_summary_op = tf.identity(train_summary_op)

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for initializer in initializers:
      initializer(sess)
    trainer_lib.restore_checkpoint(model_dir, model_saver, sess)

    num_target_iterations = restarter.target_local_step
    num_completed_iterations = num_target_iterations - restarter.steps_to_go
    with restarter.training_loop():
      for epoch_index in range(num_completed_iterations, num_target_iterations):
        summary = sess.run(train_summary_op)
        if summary_writer:
          summary_writer.add_summary(summary, epoch_index)

        if (ppo_hparams.eval_every_epochs and
            epoch_index % ppo_hparams.eval_every_epochs == 0):
          eval_summary = sess.run(eval_summary_op)
          if summary_writer:
            summary_writer.add_summary(eval_summary, epoch_index)
          if report_fn:
            summary_proto = tf.Summary()
            summary_proto.ParseFromString(eval_summary)
            for elem in summary_proto.value:
              if "mean_score" in elem.tag:
                report_fn(elem.simple_value, epoch_index)
                break

        if (model_saver and ppo_hparams.save_models_every_epochs and
            (epoch_index % ppo_hparams.save_models_every_epochs == 0 or
             (epoch_index + 1) == num_target_iterations)):
          ckpt_path = os.path.join(
              model_dir,
              "model.ckpt-{}".format(tf.train.global_step(sess, global_step))
          )
          model_saver.save(sess, ckpt_path)
          if model_save_fn:
            model_save_fn(model_dir)
def _run_train ( ppo_hparams , event_dir , model_dir , restarter , train_summary_op , eval_summary_op , initializers , report_fn = None , model_save_fn = None ) : summary_writer = tf . summary . FileWriter ( event_dir , graph = tf . get_default_graph ( ) , flush_secs = 60 ) model_saver = tf . train . Saver ( tf . global_variables ( ppo_hparams . policy_network + "/.*" ) + tf . global_variables ( "training/" + ppo_hparams . policy_network + "/.*" ) + tf . global_variables ( "global_step" ) + tf . global_variables ( "losses_avg.*" ) + tf . global_variables ( "train_stats.*" ) ) global_step = tf . train . get_or_create_global_step ( ) with tf . control_dependencies ( [ tf . assign_add ( global_step , 1 ) ] ) : train_summary_op = tf . identity ( train_summary_op ) with tf . Session ( ) as sess : sess . run ( tf . global_variables_initializer ( ) ) for initializer in initializers : initializer ( sess ) trainer_lib . restore_checkpoint ( model_dir , model_saver , sess ) num_target_iterations = restarter . target_local_step num_completed_iterations = num_target_iterations - restarter . steps_to_go with restarter . training_loop ( ) : for epoch_index in range ( num_completed_iterations , num_target_iterations ) : summary = sess . run ( train_summary_op ) if summary_writer : summary_writer . add_summary ( summary , epoch_index ) if ( ppo_hparams . eval_every_epochs and epoch_index % ppo_hparams . eval_every_epochs == 0 ) : eval_summary = sess . run ( eval_summary_op ) if summary_writer : summary_writer . add_summary ( eval_summary , epoch_index ) if report_fn : summary_proto = tf . Summary ( ) summary_proto . ParseFromString ( eval_summary ) for elem in summary_proto . value : if "mean_score" in elem . tag : report_fn ( elem . simple_value , epoch_index ) break if ( model_saver and ppo_hparams . save_models_every_epochs and ( epoch_index % ppo_hparams . save_models_every_epochs == 0 or ( epoch_index + 1 ) == num_target_iterations ) ) : ckpt_path = os . path . join ( model_dir , "model.ckpt-{}" . format ( tf . train . global_step ( sess , global_step ) ) ) model_saver . save ( sess , ckpt_path ) if model_save_fn : model_save_fn ( model_dir )
Train.