我如何将ray rllib pytorch整个模型导入到下一轮训练和随后的推理中，使用torch save load方法

Question

在ray rllib中，我通常像这样应用ray.tune.run一个ppo训练:

ray.init(log_to_driver=False, num_cpus=3, 
local_mode=args.local_mode, num_gpus=1)
env_config={"code":"codeA"}
config={
env_config={
"code":"codeA"},
"parm":"paramA"}
stop = {
"training_iteration": args.stop_iters,
"timesteps_total": args.stop_timesteps,
"episode_reward_mean": args.stop_reward,
}
results = tune.run(trainer, config=config1, verbose=0, 
stop=stop1, checkpoint_at_end=True,                               
metric='episode_reward_mean', mode="max", 
checkpoint_freq=1
)
checkpoints = results.get_trial_checkpoints_paths(
trial=results.get_best_trial(
metric='episode_reward_mean', 
mode="max"),metric='episode_reward_mean')
checkpoint_path = checkpoints[0][0]
metric = checkpoints[0][1]

在下一轮，我通常使用恢复检查点方法重新训练模型，如下所示:

results = tune.run('PPO', config=config1, verbose=0, 
stop=stop, checkpoint_at_end=True,                                   
metric='episode_reward_mean', mode="max", checkpoint_freq=1, restore=checkpoint_path)

在推理:

agent = ppo.PPOTrainer(config=config1, env=env)
agent.restore(checkpoint_path=checkpoint_path)

这些流程起作用了。问题是(1):如果我可以保存整个pytorch模型在ray.tune.run结束?(2)我可以在下一轮ray.tune.run训练中导入pytorch模型而不是恢复检查点吗?(3)在推理阶段，如何将训练好的整个pytorch模型导入PPO代理中?在还原代理推理流程中，我不能一次将10个以上的模型加载到计算机内存中。大加载显示了OOM问题。如果逐个恢复模型，则恢复检查点的过程过于耗时，无法满足时效性要求。有人能帮我吗?

Answer 1

您可以查看tune.run()中的keep_checkpoints_num和checkpoints_score_attr来自定义从这里需要多少检查点keep_checkpoints_num的默认值是None，因此它将存储所有检查点，但由于存储限制，您可以仅保留基于检查点得分属性

的最前面的检查点。