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?