由于某些原因，我需要为 Tensorflow 实现自定义资源。我试图从查找表实现中获得灵感。如果我理解得好的话，我需要实现3个TF操作：
- 创建我的资源
- 资源的初始化（例如，在查找表的情况下填充哈希表）
- 执行查找/查找/查询步骤。
为了促进实施，我依靠tensorflow/core/framework/resource_op_kernel.h。我收到以下错误
[F tensorflow/core/lib/core/refcount.h:90] Check failed: ref_.load() == 0 (1 vs. 0)
1] 29701 abort python test.py
这是重现该问题的完整代码：
using namespace tensorflow;
/** CUSTOM RESOURCE **/
class MyVector : public ResourceBase {
public:
string DebugString() override { return "MyVector"; };
private:
std::vector<int> vec_;
};
/** CREATE VECTOR **/
REGISTER_OP("CreateMyVector")
.Attr("container: string = ''")
.Attr("shared_name: string = ''")
.Output("resource: resource")
.SetIsStateful();
class MyVectorOp : public ResourceOpKernel<MyVector> {
public:
explicit MyVectorOp(OpKernelConstruction* ctx) : ResourceOpKernel(ctx) {}
private:
Status CreateResource(MyVector** resource) override {
*resource = CHECK_NOTNULL(new MyVector);
if(*resource == nullptr) {
return errors::ResourceExhausted("Failed to allocate");
}
return Status::OK();
}
Status VerifyResource(MyVector* vec) override {
return Status::OK();
}
};
REGISTER_KERNEL_BUILDER(Name("CreateMyVector").Device(DEVICE_CPU), MyVectorOp)
然后，编译后，可以使用以下 Python 代码片段重现该错误：
test_module = tf.load_op_library('./test.so')
my_vec = test_module.create_my_vector()
with tf.Session() as s:
s.run(my_vec)
作为一个附带问题，我有兴趣提供教程/指南来实现自定义资源。特别是，我想了解有关检查点/图形导出/序列化/等需要实现什么的信息。
多谢。
Add -DNDEBUG到你的构建标志。
此解决方法已解释在TF第17316期 https://github.com/tensorflow/tensorflow/issues/17316.
本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系:hwhale#tublm.com(使用前将#替换为@)
