Custom resource in Tensorflow

Asked by At

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:

  1. creation of my resource
  2. initialization of the resource (e.g. populate the hash table in case of the lookup table)
  3. implementation of the find / lookup / query step.

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

Here is the full code to reproduce the issue:

using namespace tensorflow;

class MyVector : public ResourceBase {
  string DebugString() override { return "MyVector"; };
  std::vector<int> vec_;

    .Attr("container: string = ''")
    .Attr("shared_name: string = ''")
    .Output("resource: resource")

class MyVectorOp : public ResourceOpKernel<MyVector> {
  explicit MyVectorOp(OpKernelConstruction* ctx) : ResourceOpKernel(ctx) {}

  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)

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

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

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.

Thanks a lot.

0 Answers