Tensorflow中的自定义资源

时间:2019-05-02 18:49:58

标签: c++ tensorflow tensorflow-serving

出于某些原因,我需要为Tensorflow实现自定义资源。我试图从查找表实现中获得启发。如果我理解不错,就需要实现3个TF操作:

  1. 创造我的资源
  2. 资源的初始化(例如,在查找表的情况下填充哈希表)
  3. 查找/查找/查询步骤的实现。

为了方便实施,我依靠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)

作为附带的问题,我对拥有实现自定义资源的教程/指南很感兴趣。特别是,我想了解有关检查点/图形导出/序列化等需要实现的信息。

非常感谢。

1 个答案:

答案 0 :(得分:0)

AVCaptureMovieFileOutput添加到您的构建标记中。 in TF issue 17316说明了这种解决方法。