如何在 Tensorflow C++ 中定义变量的自定义有状态 Op 保存值

How to define a custom stateful Op holding values of variables in Tensorflow C++?

本文关键字:自定义 状态 Op 保存 变量 Tensorflow C++ 定义      更新时间:2023-10-16

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

我试过REGISTER_OP((。SetIsStateful(( 但它不起作用。有谁知道如何实现这种操作?

谢谢

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

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

简而言之,定义一个包含要存储的状态的结构。并在类资源库中实现虚拟方法。然后通过查找函数还原此有状态结构,该函数从资源管理器管理的容器中检索值。

关键是您需要像这样从上下文中检索资源管理器,而不是定义一个新的资源管理器:

ResourceMgr *rm = context->resource_manager();