How to write a custom call function for a Tensorflow LSTM class?



I have defined a custom LSTM Layer as follows:

class LSTMModel(tf.keras.Model):
    def __init__(self, CNN_model, num_classes):
        self.cnn_model = CNN_model
        self.lstm = tf.keras.layers.LSTM(units=64, return_state=True, dropout=0.3)
        self.dense = tf.keras.layers.Dense(num_classes, activation="softmax")

    def call(self, input):

However, I am unclear what needs too occur in the call function here. I also wrote a generic CNN class like below:

class generic_vns_function(tf.keras.Model):
    # Where would we use layer_units here?
    def __init__(self, input_shape, layers, layer_units): 
        self.convolutions = []

        # Dynamically create Convolutional layers and MaxPools
        for layer in range(len(layers)):
            self.convolutions.append(tf.keras.layers.Conv2D(layer, 3, padding="same", 
                input_shape=input_shape, activation="relu"))

        # Add MaxPooling layer
        # Flatten 
        self.flatten = tf.keras.layers.Flatten()
        # Dense layer 
        self.dense1 = tf.keras.layers.Dense(1024, activation="relu")

    def call(self, input):
        x = input

        for layer in self.convolutions:
            x = layer(x)

        x = self.flatten(x)
        x = self.dense1(x)

        return x

but here the required structure makes a lot more sense to me. I am just initializing all of the layers. What do I need to do to initialize my LSTM layers?


You could write it like this:

import tensorflow as tf

from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import LSTM
from tensorflow.keras import Model

class LSTMModel(Model):
  def __init__(self, num_classes, num_units=64, drop_prob=0.3):
    self.num_classes = num_classes
    self.num_units = num_units
    self.drop_prob = drop_prob
    self.lstm = LSTM(
    self.dense = Dense(

  def call(self, x, training=True):
    x, *state = self.lstm(x, training=training)
    x = self.dense(x)
    return x

And then you would use it like:

model = LSTMModel(num_classes=2)

time_series = tf.random.normal((32, 64, 128))
x_pred = model(time_series)

# loss and gradients calculations ...

It is a common tensorflow idom to instantiate layers when initializing a custom layer/model, and then execute their call() methods by passing data through them in your custom call implementation.

Answered By – gobrewers14

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave A Reply

Your email address will not be published.

This website uses cookies to improve your experience. We'll assume you're ok with this, but you can opt-out if you wish. Accept Read More