python - Ray:在Python脚本中多次调用tune.run()时的内存管理

Question

我有一个 python 脚本，它使用库 ray 和 rllib 等来训练强化学习模型。该脚本使用检查点来迭代更新 rllib.PPO 模型。在每次迭代中，我都会重新定义配置并调用tune.run()，其中我将上一次迭代的检查点提供给恢复变量。对于每次调用 tune.run()，我只请求一名工作人员。在进入循环之前我初始化ray并请求大量资源。

问题在于内存快速增加，直到 ray 提示工作人员没有足够的内存并停止运行。使用 htop，我可以看到我的 python 脚本在前 10 次迭代中从 6% 上升到 80%。我想知道如何在每次迭代结束时释放资源，以便内存使用量不会随着运行时间的增加而增加。

这是我的脚本的(伪)代码:

初始化射线

ray.init(object_store_memory=50000000000，内存=50000000000)

训练循环

for iteration in range(niterations):
  new_config = ...
  prev_checkpoint = ...
  tune.run('PPO', restore= prev_checkpoint, config=new_config)    

正如你所看到的，我目前没有做任何事情来释放资源。此外，在循环内调用 ray.init() 会给出错误，表明无法多次调用它。最后，在开始时请求更多内存是不可能的，也不能解决问题，因为我想在保持内存使用量不变的情况下执行数千次迭代。

Answer 1

啊，我解决了这个问题。调用tune.run()后不需要释放任何资源，内存问题是由于在每次迭代中构建 tensorflow 图造成的。我意识到，非常烦人的是，释放由tensorflow分配的资源的唯一方法是终止python解释器(关闭tensorflow session 并不会释放它们)。因此，我编写了一个用于构建和训练该图的脚本，我将其称为“os.system()”。相当hacky，但我不知道有任何其他解决方案。