如何在 Tensorflow C++ 中定义一个自定义的有状态操作来保存变量的值？

Posted 2023-02-23

【中文标题】如何在 Tensorflow C++ 中定义一个自定义的有状态操作来保存变量的值？

【英文标题】：How to define a custom stateful Op holding values of variables in Tensorflow C++?

【发布时间】：2018-01-25 13:58:43

【问题描述】：

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

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

谢谢

【参考方案1】：

最后，我找到了如何做到这一点。

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

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

关键是你需要像这样从上下文中检索 ResourceManager 而不是定义一个新的：

ResourceMgr *rm = context->resource_manager();