calling tensorflow2 model is not returning the output that has been defined in the call() method

187 views Asked by At

When I call the tf2 model, it is not returning me the values that it is supposed to return as per the way I have defined the call() method in the tf Model subclass.

Instead, calling the call() method of the model is returning me the tensor that I have defined in the build() method

Why is this happening and how do I fix this?

import numpy as np
import tensorflow as tf

num_items = 1000
emb_dim = 32
lstm_dim = 32

class rnn_model(tf.keras.Model):
    def __init__(self, num_items, emb_dim): 
        super(rnn_model, self).__init__()
        
        self.emb   = tf.keras.layers.Embedding(num_items, emb_dim, name='embedding_layer')
        self.GRU   = tf.keras.layers.LSTM(lstm_dim, name='rnn_layer')
        self.dense = tf.keras.layers.Dense(num_items, activation = 'softmax', name='final_layer')
        
    def call(self, inp, is_training=True):  
        emb = self.emb(inp)
        gru = self.GRU(emb)
        # logits=self.dense(gru)
        
        return gru # (bs, lstm_dim=50)
    
    def build(self, inp_shape):
      x = tf.keras.Input(shape=inp_shape, name='input_layer')
      # return tf.keras.Model(inputs=[x], outputs=self.call(x))
      return tf.keras.Model(inputs=[x], outputs=self.dense(self.call(x)))

maxlen = 10
model = rnn_model(num_items, emb_dim).build((maxlen, ))
model.summary()

gru_out = model(inp)
print(gru_out.shape) # should have been (bs=16, lstm_dim=32)

Following is the output that I am getting-

Model: "functional_11"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_layer (InputLayer)     [(None, 10)]              0         
_________________________________________________________________
embedding_layer (Embedding)  (None, 10, 32)            32000     
_________________________________________________________________
rnn_layer (LSTM)             (None, 32)                8320      
_________________________________________________________________
final_layer (Dense)          (None, 1000)              33000     
=================================================================
Total params: 73,320
Trainable params: 73,320
Non-trainable params: 0
_________________________________________________________________
(16, 1000)

I intend to use only the 'final_layer' or the dense layer at the end of the model, to feed it into the sampled softmax function where it would be used with the gru_out to calculate the loss, (in order to train the model).

While testing I intend to manually pass the gru_out into model.get_layer('final_layer') to get the final logits.

0

There are 0 answers