Tensorflow 2 custom LSTM Cell initial states

Viewed 597

I try to implement a custom LSTM cell. Firstly, I try to reproduce the original LSTM cell before I put customization. However, I run into a problem that the initial state is a single tensor instead of a tuple.

class LSTMCell(keras.layers.Layer):
   def __init__(self, units, activation='tanh',
           recurrent_activation='hard_sigmoid',
           use_bias=True,
           kernel_initializer='glorot_uniform',
           recurrent_initializer='orthogonal',
           bias_initializer='zeros',**kwargs):
    self.units = units
    self.state_size = units
    self.kernel_initializer = kernel_initializer
    self.use_bias = use_bias
    self. recurrent_initializer = recurrent_initializer
    self.bias_initializer = bias_initializer
    super(LSTMCell, self).__init__(**kwargs)
def build(self, input_shape):
    input_dim = input_shape[-1]
    self.kernel = self.add_weight(
        shape=(input_dim, self.units * 4),
        name='kernel',
        initializer=self.kernel_initializer)
    self.recurrent_kernel = self.add_weight(
        shape=(self.units, self.units * 4),
        name='recurrent_kernel',
        initializer=self.recurrent_initializer)
    self.bias = self.add_weight(
      shape=(self.units * 4,),
      name='bias',
      initializer=self.bias_initializer)
    
    
       
def _compute_carry_and_output_fused(self, z, c_tm1):
    z0, z1, z2, z3 = z
    i = K.sigmoid(z0)
    f = K.sigmoid(z1)
    c = f * c_tm1 + i * K.tanh(z2)
    o = K.sigmoid(z3)
    return c, o
    
def call(self, inputs, states, training=None):
    
    
   
    h_tm1 = states[0]  # previous memory state
    c_tm1 = states[1]  # previous carry state
    
    z = K.dot(inputs, self.kernel)
    z += K.dot(h_tm1, self.recurrent_kernel)
    z = K.bias_add(z, self.bias)
    z = tf.split(z, num_or_size_splits=4, axis=1)
    c, o = self._compute_carry_and_output_fused(z, c_tm1)

    h = o * K.sigmoid(c)
    self.h = h
    self.c = c
    
    return h, [h,c]

cell= LSTMCell(32)  
layer = RNN(cell) 
a = np.random.rand(44,10,40)
out = layer(a)

I got error message: c_tm1 = states[1]

IndexError: tuple index out of range

I guess the problem is that the get_init_state function of a plain rnn cell was called, which returned a single tensor. I tried to solve this issue by initializing the states inside my class:

self.initial = True
.....

if self.initial:
 h_tm1 = tf.zeros(shape=[inputs.shape[0], self.state_size], name='h')
 c_tm1 = tf.zeros(shape=[inputs.shape[0], self.state_size], name='c')
 self.initial = False

'''

But it did not work either. How can I make my LSTMCell work by letting it call the right get_init_state function or in whatever way?

2 Answers

I just tested your code with Tensorflow 2.7.0 package and seems like your code is working just fine.

Perhaps your defined functions (def __ init __, def build, def _compute_carry_and_output_fused, and def call) are not properly aligned insde the class LSTMCell.

And, perhaps your self instance objects (units, state_size, ... , bias_initializer) and super init are also not aligned inside the init defined function.

You should give proper shape with self.state_size:

self.state_size = (self.units, self.units)
Related