Как написать пользовательскую операцию Tensorflow, содержащую постоянный объект C ++? - PullRequest
Новая
       83

Как написать пользовательскую операцию Tensorflow, содержащую постоянный объект C ++?

3 голосов
/ 30 мая 2019

Я разрабатываю модель последовательности Tensorflow, которая использует поиск луча по графу декодирования OpenFST (загружается из двоичного файла) по логитам, выводимым из модели последовательности Tensorflow.

Я написал пользовательскийОператор, который позволяет мне выполнять декодирование по логитам, но каждый раз перед операцией декодирования у меня вызывается операционный вызов fst :: Read (BINARY_FILE).Это может быть хорошо, пока оно остается маленьким, но я бы хотел избежать накладных расходов на ввод / вывод.

Я прочитал пользовательскую операцию 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);

  }
};

В питоне: 


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

1 Ответ

1 голос
/ 31 мая 2019

Решение:

Не понимал, что вы можете установить личные атрибуты, используя "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);

  }
};

В питоне: 


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