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.