python - 如何在 Tensorflow C++ 中定义自定义有状态 Op 变量值？

Question

我试图用状态定义一个新的 Op，然后用 python 包装它。问题是我希望这个操作在 Session.run() 之间保持变量(就像在这个操作中初始化一个 int foo 然后每次运行时将 foo 添加 1。预期的是 foo 比上次运行多 1。)

我已经尝试过 REGISTER_OP().SetIsStateful() 但它不起作用。有谁知道如何实现这种操作？

谢谢

Answer 1

终于，我找到了如何做到这一点。

根据tensorflow变量的实现。状态值可以存储在定义的资源管理器中here .

简而言之，定义一个包含您要存储的状态的结构。并在 Class ResourceBase 中实现虚方法。然后通过 Lookup 函数恢复这个有状态结构，该函数从 Resource Manager 管理的容器中检索值。

重点是您需要像这样从上下文中检索 ResourceManager 而不是定义一个新的 ResourceManager:

ResourceMgr *rm = context->resource_manager();