如何编写包含持久 C++ 对象的 Tensorflow 自定义操作？

Question

我正在开发一个 Tensorflow 序列模型，该模型通过 OpenFST 解码图（从二进制文件加载）对 Tensorflow 序列模型的 logits 输出使用波束搜索。

我编写了一个自定义操作，允许我对 logits 执行解码，但每次，我都会在执行解码之前调用操作 fst::Read(BINARY_FILE)。只要它保持很小，这可能没问题，但我想避免 I/O 开销。

我通读了 Tensorflow 自定义操作并试图找到类似的示例，但我仍然迷路了。基本上，我想在图中做的是：

FstDecodingOp.Initialize('BINARY_FILE.bin') #loads the BINARY_FILE.bin into memory
...
for o in output:
    FstDecodingOp.decode(o) # uses BINARY_FILE.bin to decode

这在张量流图之外的 Python 中当然会很简单，但我最终需要将其移至普通 TF-Serving 环境中，因此需要将其冻结到导出图中。有没有人遇到过类似的问题？

解法：

没有意识到您可以使用 "OpKernel(context)" 设置私有属性。刚刚使用该函数对其进行了初始化。

编辑：关于我是如何做到的更多细节。还没有尝试服务。

REGISTER_OP("FstDecoder")
    .Input("log_likelihoods: float")
    .Attr("fst_decoder_path: string")
    ....

...

template <typename Device, typename T>
class FstDecoderOp : public OpKernel {

private:
   fst::Fst<fst::StdArc>* fst_;
   float beam_;

public:
  explicit FstDecoderOp(OpKernelConstruction* context) : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("beam", &beam_));

    std::string fst_path;
    OP_REQUIRES_OK(context, context->GetAttr("fst_decoder_path", &fst_path));

    fst_ = fst::Fst<fst::StdArc>::Read(fst_path);
  }

  void Compute(OpKernelContext* context) override {
    // do some compute 
    const Tensor* log_likelihoods;

    OP_REQUIRES_OK(context, context->input("log_likelihoods", 
     &log_likelihoods));

    // simplified 
    compute_op(_fst, log_likelihoods);

  }
};

在python中：


sess = tf.Session()
mat = tf.placeholder(tf.float32, shape=test_npy.shape)
res_ = decoder_op.fst_decoder(beam=30, fst_decoder_path="decoder_path.fst", log_likelihoods=mat)
res = sess.run(res_, {mat : test_npy} )

Answer 1

解法：

没有意识到您可以使用 "OpKernel(context)" 设置私有属性。刚刚使用该函数对其进行了初始化。

编辑：关于我是如何做到的更多细节。还没有尝试服务。

REGISTER_OP("FstDecoder")
    .Input("log_likelihoods: float")
    .Attr("fst_decoder_path: string")
    ....

...

template <typename Device, typename T>
class FstDecoderOp : public OpKernel {

private:
   fst::Fst<fst::StdArc>* fst_;
   float beam_;

public:
  explicit FstDecoderOp(OpKernelConstruction* context) : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("beam", &beam_));

    std::string fst_path;
    OP_REQUIRES_OK(context, context->GetAttr("fst_decoder_path", &fst_path));

    fst_ = fst::Fst<fst::StdArc>::Read(fst_path);
  }

  void Compute(OpKernelContext* context) override {
    // do some compute 
    const Tensor* log_likelihoods;

    OP_REQUIRES_OK(context, context->input("log_likelihoods", 
     &log_likelihoods));

    // simplified 
    compute_op(_fst, log_likelihoods);

  }
};

在python中：


sess = tf.Session()
mat = tf.placeholder(tf.float32, shape=test_npy.shape)
res_ = decoder_op.fst_decoder(beam=30, fst_decoder_path="decoder_path.fst", log_likelihoods=mat)
res = sess.run(res_, {mat : test_npy} )

如何编写包含持久 C++ 对象的 Tensorflow 自定义操作？

How do I write a Tensorflow custom op containing a persistent C++ object?

python

nlp

openfst

tensorflow

tensorflow-serving