I am trying to change the training and inference signature functions of my model that is going to be exported as a .tflite model to be used on a smartphone for on-device training. As the model is now, the body of it is a MobileNetV2 and the head which is the part that is trainable on the device is just a fully-connected layer with softMax. I am trying to alter the training and inference functions so that the last 2 layers of the mobileNetV2 are being used again on the head as 2 hidden layers before my softMax activations of the fully-connected layer.
This is my initialization where I tried to save the second to last and last layer:
def __init__(self, learning_rate=0.001):
"""Initializes a transfer learning model instance.
Args:
learning_rate: A learning rate for the optimzer.
"""
self.num_features = NUM_FEATURES
self.num_classes = NUM_CLASSES
# trainable weights and bias for softmax
self.ws = tf.Variable(
tf.zeros((self.num_features, self.num_classes)),
name='ws',
trainable=True)
self.bs = tf.Variable(
tf.zeros((1, self.num_classes)), name='bs', trainable=True)
# base model
self.base = tf.keras.applications.MobileNetV2(
input_shape=(IMG_SIZE, IMG_SIZE, 3),
alpha=1.0,
include_top=False,
weights='imagenet')
# loss function and optimizer
self.loss_fn = tf.keras.losses.CategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
# save the last layer of the base model
self.pre_last_layer = self.base.layers[-2].output
self.last_layer = self.base.layers[-1].output
and this is my altered training function where I added those 2 layers.
@tf.function(input_signature=[
tf.TensorSpec([None, NUM_FEATURES], tf.float32),
tf.TensorSpec([None, NUM_CLASSES], tf.float32),
])
def train(self, bottleneck, label):
"""Runs one training step with the given bottleneck features and labels.
Args:
bottleneck: A tensor of bottleneck features generated from the base model.
label: A tensor of class labels for the given batch.
Returns:
Map of the training loss.
"""
with tf.GradientTape() as tape:
x = tf.keras.layers.Reshape((7, 7, 1280))(bottleneck)
x = self.pre_last_layer(x, training = True)
x = self.last_layer(x, training=True)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
logits = tf.matmul(x, self.ws) + self.bs
prediction = tf.nn.softmax(logits)
loss = self.loss_fn(prediction, label)
gradients = tape.gradient(loss, [self.ws, self.bs])
self.optimizer.apply_gradients(zip(gradients, [self.ws, self.bs]))
result = {'loss': loss}
for grad in gradients:
result[grad.name] = grad
return result
As it is now, I keep getting TypeError: 'KerasTensor' object is not callable where im using
x = self.pre_last_layer(x, training= True) and I cant find a way to fix it.
Any help is appreciated. Just a heads up that I'm new to tensorflow and tensorflow-lite. Cheers!
Searching online the TypeError, I found that it's caused if you try to pass an input to an input tensor which I didn't identify happening here at least to my knowledge.
