Ray: [rllib]How to load the pre-trained PPO model with a lower overhead?

[xudongliao]

How to load pre-trained model with lower overhead?

• Ray 0.8.0
• Tensorflow 1.14.0
• OS: Ubuntu 18.04
• Python: 3.5.6 :: Anaconda

I trained a PPO1 model with:

def train():
    ray.init()
    config = ppo.DEFAULT_CONFIG.copy()
    config["num_gpus"] = 0
    # config["num_workers"] = int(multiprocessing.cpu_count() / 2)
    config["num_workers"] = 4
    config["eager"] = False
    config["log_level"] = "INFO"
    config["monitor"] = True
    config["num_cpus_per_worker"] = 0
    trainer = ppo.PPOTrainer(config=config, env=ShimNetworkEnv)

    for i in range(1000):
        # Perform one iteration of training the policy with PPO
        result = trainer.train()
        print(pretty_print(result))

        if i % 100 == 0:
            checkpoint = trainer.save()
            print("checkpoint saved at", checkpoint)

And I'd like to use the checkpoint for further testing with:

class Model(object):
    def __init__(self, ckpt_path):
        ray.init()
        config = ppo.DEFAULT_CONFIG.copy()
        config['num_workers'] = 1

        agent = ppo.PPOTrainer(config, env=ShimNetworkEnv)
        agent.restore(ckpt_path)
        self.policy = agent.workers.local_worker().get_policy()

    def act(self, state):
        action = self.policy.compute_actions([state])
        return action[0][0]

But loading the model in this way would incur an annoying overhead (> 5s) on creating Trainer object and restoring the checkpoint. It's unacceptable in my working environment. Thus is there any ways to lower this loading overhead? Thanks.

[ericl]

Try config['num_workers'] = 0 which should avoid creating any actors at all.

Other than that, the only faster way is to export the model for TF serving: https://ray.readthedocs.io/en/latest/rllib-package-ref.html#ray.rllib.policy.Policy.export_model

[xudongliao]

This solution works well. Thanks :)

[stale[bot]]

Hi, I'm a bot from the Ray team :)

To help human contributors to focus on more relevant issues, I will automatically add the stale label to issues that have had no activity for more than 4 months.

If there is no further activity in the 14 days, the issue will be closed!

• If you'd like to keep the issue open, just leave any comment, and the stale label will be removed!
• If you'd like to get more attention to the issue, please tag one of Ray's contributors.

You can always ask for help on our discussion forum or Ray's public slack channel.