title: How do I write a Tensorflow custom op containing a persistent C++ object?
body: <p>I'm developing a Tensorflow sequence model that uses a beam search through an OpenFST decoding graph (loaded from a binary file) over the logits output from a Tensorflow sequence model.</p>
  
  <p>I've written a custom op that allows me to perform decoding over the logits, but each time, I'm having the op call fst::Read(BINARY_FILE) before performing the decoding. This might be fine as long as it stays small but I'd like to avoid the I/O overhead.</p>
  
  <p>I've read through the Tensorflow custom op and tried to find similar examples but I'm still lost. Basically, what I want to do in the graph is:</p>
  
  <pre class="lang-py prettyprint-override"><code>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
  
  </code></pre>
  
  <p>This would of course be straightforward in Python outside of the tensorflow graph, but I need to eventually move this into a vanilla TF-Serving environment, so it needs to get frozen into an export graph. Has anyone encountered a similar problem before? </p>
  
  <p><strong>Solution:</strong></p>
  
  <p>Didn't realize that you could set private attributes using the "OpKernel(context)". Just initialized it using that function.</p>
  
  <p><strong>Edit: more detail on how I did it. Have yet to try serving.</strong></p>
  
  <pre class="lang-cpp prettyprint-override"><code>REGISTER_OP("FstDecoder")
      .Input("log_likelihoods: float")
      .Attr("fst_decoder_path: string")
      ....
  
  ...
  
  template &lt;typename Device, typename T&gt;
  class FstDecoderOp : public OpKernel {
  
  private:
     fst::Fst&lt;fst::StdArc&gt;* fst_;
     float beam_;
  
  public:
    explicit FstDecoderOp(OpKernelConstruction* context) : OpKernel(context) {
      OP_REQUIRES_OK(context, context-&gt;GetAttr("beam", &amp;beam_));
  
      std::string fst_path;
      OP_REQUIRES_OK(context, context-&gt;GetAttr("fst_decoder_path", &amp;fst_path));
  
      fst_ = fst::Fst&lt;fst::StdArc&gt;::Read(fst_path);
    }
  
    void Compute(OpKernelContext* context) override {
      // do some compute 
      const Tensor* log_likelihoods;
  
      OP_REQUIRES_OK(context, context-&gt;input("log_likelihoods", 
       &amp;log_likelihoods));
  
      // simplified 
      compute_op(_fst, log_likelihoods);
  
    }
  };
  </code></pre>
  
  <p>In python:</p>
  
  <pre class="lang-py prettyprint-override"><code>
  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} )
  
  </code></pre>