def train()

in lm_human_preferences/train_policy.py [0:0]

• 81 lines of code
• 14 McCabe index (conditional complexity)

def train(hparams: HParams):
    save_dir = hparams.run.save_dir
    if hparams.rewards.train_new_model:
        assert hparams.task == hparams.rewards.train_new_model.task, f'{hparams.task} != {hparams.rewards.train_new_model.task}'
        hparams.rewards.train_new_model.run.save_dir = save_dir
        train_reward.train(hparams.rewards.train_new_model)
        if 'pytest' in sys.modules:
            hparams.rewards.trained_model = 'test'
        elif save_dir:
            hparams.rewards.trained_model = None if save_dir is None else os.path.join(save_dir, 'reward_model')

    comm = MPI.COMM_WORLD

    with tf.Graph().as_default():
        hyperparams.dump(hparams)

        m = trained_models.TrainedModel(hparams.task.policy.initial_model)
        encoder = m.encoding.get_encoder()
        hyperparams.dump(m.hparams(), name='model_hparams')

        if save_dir:
            if not save_dir.startswith('https:'):
                os.makedirs(os.path.join(save_dir, 'policy'), exist_ok=True)
            with tf.gfile.Open(os.path.join(save_dir, 'train_policy_hparams.json'), 'w') as f:
                json.dump(hparams.to_nested_dict(), f, indent=2)
            with tf.gfile.Open(os.path.join(save_dir, 'policy', 'hparams.json'), 'w') as f:
                json.dump(m.hparams().to_nested_dict(), f, indent=2)
            with tf.gfile.Open(os.path.join(save_dir, 'policy', 'encoding'), 'w') as f:
                json.dump(m.encoding.name, f, indent=2)
        utils.set_mpi_seed(hparams.run.seed)

        score_model = TrainedRewardModel(hparams.rewards.trained_model, m.encoding, comm=comm)

        ref_policy = Policy(
            m, scope='ref_policy',
            is_root=comm.Get_rank() == 0,
            embed_queries=lm_tasks.query_formatter(hparams.task, encoder),
            temperature=hparams.task.policy.temperature,
            build_respond=False)

        policy = Policy(
            m, scope='policy',
            is_root=comm.Get_rank() == 0,
            embed_queries=lm_tasks.query_formatter(hparams.task, encoder),
            temperature=hparams.task.policy.temperature)

        query_sampler = lm_tasks.make_query_sampler(
            hparams=hparams.task, encoder=encoder, comm=comm,
            batch_size=utils.exact_div(hparams.ppo.batch_size, comm.Get_size()),
        )

        per_rank_minibatch_size = utils.exact_div(hparams.ppo.batch_size, hparams.ppo.nminibatches * comm.Get_size())
        if hparams.ppo.whiten_rewards:
            assert per_rank_minibatch_size >= 8, \
                f"Per-rank minibatch size {per_rank_minibatch_size} is insufficient for whitening"

        global_step = tf.train.get_or_create_global_step()
        increment_global_step = tf.group(global_step.assign_add(1))

        with utils.variables_on_gpu():

            ppo_trainer = PPOTrainer(
                policy=policy, ref_policy=ref_policy, query_sampler=query_sampler,
                score_fn=make_score_fn(hparams.task, score_model=score_model),
                hparams=hparams, comm=comm)

        if comm.Get_rank() == 0 and save_dir:
            print(f"Will save to {save_dir}")
            saver = tf.train.Saver(max_to_keep=20, save_relative_paths=True)
            checkpoint_dir = os.path.join(save_dir, 'policy/checkpoints/model.ckpt')
        else:
            saver = None
            checkpoint_dir = None

        @utils.graph_function()
        def sync_models():
            score_model.ensure_built()
            return utils.variable_synchronizer(comm, vars=score_model.get_params() + ref_policy.get_params() + policy.get_params())

        init_ops = tf.group(
            tf.global_variables_initializer(),
            tf.local_variables_initializer(),
            summary.summary_writer_initializer_op())

        with utils.mpi_session() as sess:
            init_ops.run()

            sync_models()

            tf.get_default_graph().finalize()

            try:
                while global_step.eval() < nupdates(hparams):
                    ppo_trainer.step()
                    increment_global_step.run()

                    if saver and global_step.eval() % hparams.run.save_interval == 0:
                        saver.save(sess, checkpoint_dir, global_step=global_step)
            finally:
                if saver:
                    saver.save(sess, checkpoint_dir, global_step=global_step)