tastyminerals
tastyminerals

Reputation: 6538

Is it possible to create custom Tensorflow GRU/LSTM cell?

I'd like to modify the tf.nn.rnn_cell.GRUCell and add another gate to it for second input such that in addition to z – update gate and r – reset gate there would be third g - custom gate for the second input to the network similarly to Kiddon et al. 2016 and Wen et al. 2015.

enter image description here

I managed to do this for Torch7 GRU cell and it successfully worked for the specific sequence modeling problem. I suspect this is not straightforward to implement in Tensorflow but maybe I am wrong. Can somebody please provide more details on custom cells in Tensorflow?

Upvotes: 3

Views: 1610

Answers (1)

Jeffrey Ede
Jeffrey Ede

Reputation: 631

It is possible. As illustrated by this article, LSTM and GRU cells are just arrangements of non-linearities and arithmetic operations.

This repository contains custom LSTM, GRU and other RNN cell implementations for pyTorch.

This repository contains custom LSTM and GRU implementations for TensorFlow.

Upvotes: 1

Related Questions