如何编写自定义C ++操作来更新资源变量?

时间:2019-04-13 01:31:10

标签: c++ tensorflow

我正在尝试更新诸如https://github.com/openai/blocksparse/blob/master/blocksparse/optimize.py中的优化器,以支持存储为资源变量的参数。查看内置优化器ops的TensorFlow source code,似乎有必要使用相对复杂的锁定和写时复制逻辑来处理资源变量可能处于的不同状态。TensorFlow在{ {3}}为此,调用MaybeLockVariableInputMutexesInOrderGetInputTensorFromVariable

不幸的是,training_op_helpers.h不在公共TensorFlow标头中,所以我不能只是#include它。更糟糕的是,它依赖于libtensorflow_framework.so甚至不导出的内部功能(特别是tensorflow/core/kernels/training_op_helpers.htensorflow::functor::DenseUpdate),因此复制标头也不足够。

是否有一种预定的方式来做这种事情?在这一点上,我很想编译自己的TensorFlow分支,以导出这些功能。 (我想重新实现它们不会太糟糕)。

0 个答案:

没有答案