want zoneout lstm supported

[breadbread1984]

Describe the feature and the current behavior/state.

no zoneout lstm found in addons

Relevant information

• Are you willing to contribute it (yes/no): no
• Are you willing to maintain it going forward? (yes/no): no
• Is there a relevant academic paper? (if so, where): https://arxiv.org/abs/1606.01305
• Is there already an implementation in another framework? (if so, where): https://github.com/LaurentMazare/deep-models/blob/c54a0f5227d1b03018bbbf19075454250266bf90/rhn/lstm_zoneout.py
• Was it part of tf.contrib? (if so, where):

Which API type would this fall under (layer, metric, optimizer, etc.)

tfa.layers

Who will benefit with this feature?

people implementing Tacotron2 with tf.keras API

Any other info.

[failure-to-thrive]

I could handle it...

[breadbread1984]

really appreciated to it

[failure-to-thrive]

So far so good! Any ideas on how to test it? Input values, random seed, expected output values?

[breadbread1984]

speedy implement! given the same input and zero initial state, it outputs a tensor having element value either the same as lstm's or zero.

[failure-to-thrive]

Here it is! You can copy & paste the class into your code. Example of usage and tests are further down below. If everything work as intended we could incorporate it into TFA.

import tensorflow as tf
from tensorflow.keras.layers import LSTMCell, RNN, LSTM

class ZoneoutLSTMCell(LSTMCell):
    def __init__(
        self,
        units,
        zoneout_h = 0,
        zoneout_c = 0,
        **kwargs
    ):
        super().__init__(
            units,
            **kwargs
        )
        self.zoneout_h = zoneout_h
        self.zoneout_c = zoneout_c

    def _zoneout(self, t, tm1, rate, training):
        dt = tf.cast(tf.random.uniform(t.shape) >= rate * training, t.dtype)
        return dt*t + (1 - dt)*tm1

    def call(self, inputs, states, training=None):
        output, new_states = super().call(inputs, states, training)
        h = self._zoneout(new_states[0], states[0], self.zoneout_h, training)
        c = self._zoneout(new_states[1], states[1], self.zoneout_c, training)
        return h, [h, c]


x = tf.constant([[[1., 2, 3, 4, 5]]])
initial_state = [tf.constant([[11., 12, 13]]), tf.constant([[14., 15, 16]])]

tf.random.set_seed(0)
l0 = LSTM(3, return_state=True)
y0 = l0(x, initial_state=initial_state, training=True)
tf.print(y0)

tf.random.set_seed(0)
l = RNN(ZoneoutLSTMCell(3, zoneout_h=.3, zoneout_c=.5), return_state=True)
y = l(x, initial_state=initial_state, training=True)
tf.print(y)

[breadbread1984]

you need to save zoneout_h and zoneout_c in operator's config dictionary.

[breadbread1984]

I get unsupported operand type(s) for *: 'float' and 'NoteType' at line "dt = tf.cast(tf.random.uniform(t.shape) >= rate * training, t.dtype)"

[failure-to-thrive]

you need to save zoneout_h and zoneout_c in operator's config dictionary.

Sure. The code above is just an algorithm implementation. If it is OK we can move forward.

I get unsupported operand type(s) for *: 'float' and 'NoteType' at line "dt = tf.cast(tf.random.uniform(t.shape) >= rate * training, t.dtype)"

Could you describe environment where it happened?

[breadbread1984]

sorry, I used the operator the wrong way. I can successfully run the code.