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

Question

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

我编写了一个自定义操作，允许我对 logit 执行解码，但每次，我在执行解码之前都会调用 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-Serve 环境中，因此需要将其冻结到导出图中。以前有没有人遇到过类似的问题？

溶液：

没有意识到您可以使用"OpKernel(上下文)"设置私有属性。只是使用该函数初始化了它。

编辑：关于我如何做到这一点的更多详细信息。尚未尝试食用。

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);
}
};

在蟒蛇中：

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(上下文)"设置私有属性。只是使用该函数初始化了它。

编辑：关于我如何做到这一点的更多详细信息。尚未尝试食用。

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);
}
};

在蟒蛇中：

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} )