Пользовательский ресурс в Tensorflow - PullRequest
1 голос
/ 02 мая 2019

По некоторым причинам мне нужно реализовать собственный ресурс для 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)

В качестве дополнительного вопроса мне было бы интересно иметь учебные пособия / руководства по реализации пользовательских ресурсов.В частности, я хотел бы получить информацию о том, что необходимо реализовать для контрольных точек / экспорта графиков / сериализации / и т. Д.

Большое спасибо.

Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...