I am trying to implement a simple LSTMCell without the “fancy kwargs” defaultly implemented in the tf.keras.layers.LSTMCell class, following a schematic model like this. It doesn’t really have a direct purpose, I would just like to practice implementing a more complex RNNCell than the one described here in the Examples section. My code is the following:
from keras import Inputfrom keras.layers import Layer, RNNfrom keras.models import Modelimport keras.backend as Kclass CustomLSTMCell(Layer): def __init__(self, units, **kwargs): self.state_size = units super(CustomLSTMCell, self).__init__(**kwargs) def build(self, input_shape): self.forget_w = self.add_weight(shape=(self.state_size, self.state_size + input_shape[-1]), initializer='uniform', name='forget_w') self.forget_b = self.add_weight(shape=(self.state_size,), initializer='uniform', name='forget_b') self.input_w1 = self.add_weight(shape=(self.state_size, self.state_size + input_shape[-1]), initializer='uniform', name='input_w1') self.input_b1 = self.add_weight(shape=(self.state_size,), initializer='uniform', name='input_b1') self.input_w2 = self.add_weight(shape=(self.state_size, self.state_size + input_shape[-1]), initializer='uniform', name='input_w2') self.input_b2 = self.add_weight(shape=(self.state_size,), initializer='uniform', name='input_b2') self.output_w = self.add_weight(shape=(self.state_size, self.state_size + input_shape[-1]), initializer='uniform', name='output_w') self.output_b = self.add_weight(shape=(self.state_size,), initializer='uniform', name='output_b') self.built = True def merge_with_state(self, inputs): self.stateH = K.concatenate([self.stateH, inputs], axis=-1) def forget_gate(self): forget = K.dot(self.forget_w, self.stateH) + self.forget_b forget = K.sigmoid(forget) self.stateC = self.stateC * forget def input_gate(self): candidate = K.dot(self.input_w1, self.stateH) + self.input_b1 candidate = K.tanh(candidate) amount = K.dot(self.input_w2, self.stateH) + self.input_b2 amount = K.tanh(amount) self.stateC = self.stateC + amount * candidate def output_gate(self): self.stateH = K.dot(self.output_w, self.stateH) + self.output_b self.stateH = K.sigmoid(self.stateH) self.stateH = self.stateH * K.tanh(self.stateC) def call(self, inputs, states): self.stateH = states[0] self.stateC = states[1] self.merge_with_state(inputs) self.forget_gate() self.input_gate() self.output_gate() return self.stateH, [self.stateH, self.stateC]# Testinginp = Input(shape=(None, 3))lstm = RNN(CustomLSTMCell(10))(inp)model = Model(inputs=inp, outputs=lstm)inp_value = [[[[1,2,3], [2,3,4], [3,4,5]]]]pred = model.predict(inp_value)print(pred)However, when I tried to test it an exception was raised with the following message:
IndexError: tuple index out of rangein the call function at the line where I set the value for self.stateC. Here, I thought that initially the statesargument of the call function is a tensor and not a list of tensors, so this is why I get an error. So I added a self.already_called = False line to the classes __init__ and the following segment to the call function:
if not self.already_called: self.stateH = K.ones(self.state_size) self.stateC = K.ones(self.state_size) self.already_called = True else: self.stateH = states[0] self.stateC = states[1]hoping that it will eliminate the problem. This resulted in another error at the merge_with_state function:
ValueError: Shape must be rank 1 but is rank 2 for 'rnn_1/concat' (op: 'ConcatV2') with input shapes: [10], [?,3], [].which I genuinely do not get, since the RNN layer should only “show” the CustomLSTMCell tensors with shape (3) and not (None, 3), since axis 0 is the axis it should iterate along. At this point I was convinced that I am doing something really wrong and should ask the community for help. Basically my question is: what is wrong with my code and if “almost everything”, then how should I implement an LSTMCell from scratch?
Advertisement
Answer
Ok, so it seems that I managed to fix the problem. It turns out that it is always useful to read the documentation, in this case the docs for the RNN class. First, the already_called attribute is unnecessary, because the problem lies in the first line of the __init__ function: the state_size attribute should be a list of integers and not only one integer, like this: self.state_size = [units, units] (since we need two states for an LSTM of size units and not one). When I corrected it I got a different error: the tensors are not compatible in dimension in the forget_gate for addition. This happened because the RNN sees the whole batch at once and not each element in the batch separately (thus the None shape at axis 0). The correction for it is to add an extra dimension to each tensor of size 1 at axis 0 like this:
self.forget_w = self.add_weight(shape=(1, self.state_size, self.state_size + input_shape[-1]), initializer='uniform', name='forget_w')and instead of dot products I had to use the K.batch_dot function. So the whole, working code is the following:
from keras import Input from keras.layers import Layer, RNN from keras.models import Model import keras.backend as K class CustomLSTMCell(Layer): def __init__(self, units, **kwargs): self.state_size = [units, units] super(CustomLSTMCell, self).__init__(**kwargs) def build(self, input_shape): self.forget_w = self.add_weight(shape=(1, self.state_size[0], self.state_size[0] + input_shape[-1]), initializer='uniform', name='forget_w') self.forget_b = self.add_weight(shape=(1, self.state_size[0]), initializer='uniform', name='forget_b') self.input_w1 = self.add_weight(shape=(1, self.state_size[0], self.state_size[0] + input_shape[-1]), initializer='uniform', name='input_w1') self.input_b1 = self.add_weight(shape=(1, self.state_size[0]), initializer='uniform', name='input_b1') self.input_w2 = self.add_weight(shape=(1, self.state_size[0], self.state_size[0] + input_shape[-1]), initializer='uniform', name='input_w2') self.input_b2 = self.add_weight(shape=(1, self.state_size[0],), initializer='uniform', name='input_b2') self.output_w = self.add_weight(shape=(1, self.state_size[0], self.state_size[0] + input_shape[-1]), initializer='uniform', name='output_w') self.output_b = self.add_weight(shape=(1, self.state_size[0],), initializer='uniform', name='output_b') self.built = True def merge_with_state(self, inputs): self.stateH = K.concatenate([self.stateH, inputs], axis=-1) def forget_gate(self): forget = K.batch_dot(self.forget_w, self.stateH) + self.forget_b forget = K.sigmoid(forget) self.stateC = self.stateC * forget def input_gate(self): candidate = K.batch_dot(self.input_w1, self.stateH) + self.input_b1 candidate = K.tanh(candidate) amount = K.batch_dot(self.input_w2, self.stateH) + self.input_b2 amount = K.sigmoid(amount) self.stateC = self.stateC + amount * candidate def output_gate(self): self.stateH = K.batch_dot(self.output_w, self.stateH) + self.output_b self.stateH = K.sigmoid(self.stateH) self.stateH = self.stateH * K.tanh(self.stateC) def call(self, inputs, states): self.stateH = states[0] self.stateC = states[1] self.merge_with_state(inputs) self.forget_gate() self.input_gate() self.output_gate() return self.stateH, [self.stateH, self.stateC] inp = Input(shape=(None, 3)) lstm = RNN(CustomLSTMCell(10))(inp) model = Model(inputs=inp, outputs=lstm) inp_value = [[[[1,2,3], [2,3,4], [3,4,5]]]] pred = model.predict(inp_value) print(pred)Edit: In the question I made a mistake with respect to the model linked and used a tanh function in the input_gate for amount instead of a sigmoid. Here I edited it in the code, so it is correct now.