Tensorflow 中的自定义资源

本文介绍了Tensorflow 中的自定义资源的处理方法，对大家解决问题具有一定的参考价值，需要的朋友们下面随着小编来一起学习吧！

问题描述

出于某些原因，我需要为 Tensorflow 实现自定义资源.我试图从查找表的实现中获得灵感.如果我理解的好，我需要实现3个TF操作:

For some reasons, I need to implement a custom resource for Tensorflow. I tried to get inspiration from lookup table implementations. If I understood well, I need to implement 3 TF operations:

创建我的资源
资源的初始化(例如，在查找表的情况下填充哈希表)
执行查找/查找/查询步骤.

为了便于实现，我依赖于 tensorflow/core/framework/resource_op_kernel.h.我收到以下错误

To facilitate the implementation, I'm relying on tensorflow/core/framework/resource_op_kernel.h. I get the following error

[F tensorflow/core/lib/core/refcount.h:90] Check failed: ref_.load() == 0 (1 vs. 0)
1]    29701 abort      python test.py

这是重现问题的完整代码:

Here is the full code to reproduce the issue:

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 代码片段重现:

and then, after compiling, the error can be reproduced with this Python snippet of code:

test_module = tf.load_op_library('./test.so')
my_vec = test_module.create_my_vector()
with tf.Session() as s:
  s.run(my_vec)

作为一个附带问题，我有兴趣获得实施自定义资源的教程/指南.特别是，我想了解有关检查点/图形导出/序列化等需要实现的内容的信息.

As a side question, I'd be interested in having tutorials / guidelines to implement custom resources. In particular, I'd like to have information about what needs to be implemented for checkpoints / graph export / serialization / etc.

非常感谢.

17316

Tensorflow 中的自定义资源

问题描述

推荐答案