[Python] Keras - Replace intermediate layer in Functional model

I am trying to make a custom implementation of LoRA for Keras, and I am running into some trouble trying to get my approach to work. I am using Keras 3.4.1 inside Google Collab.

First, I have tried the solution explained in this Stack Overflow post and it works for me when I use VGG16. However, I am trying to come up with a solution that would work for Functional models, not only for sequential models such as VGG16.

Therefore, I am trying to do something similar to what is explained in this Keras tutorial, but it is not working for me. First, I create the model without the last layer, disable training on it and print its summary:

vgg16_lora = tf.keras.applications.vgg16.VGG16(weights="imagenet", input_shape=(224, 224, 3))
vgg16_lora.trainable = False

lora_model = tf.keras.Model(inputs=vgg16_lora.inputs, outputs=vgg16_lora.layers[-2].output)
lora_model.summary()


Output:

Then I try to add the LoRA layer as follows:

import math

class LoraLayer(tf.keras.Layer):
def __init__(self, original_layer, rank=8, trainable=False, **kwargs):
original_layer_config = original_layer.get_config()
name = original_layer_config["name"]
kwargs.pop('name', None)

super().__init__(name=name, trainable=trainable, **kwargs)

self.original_layer = original_layer
self.original_layer.trainable = False

self.A = tf.keras.layers.Dense(units=rank, use_bias=False, kernel_initializer=tf.keras.initializers.VarianceScaling(scale=math.sqrt(5), mode="fan_in", distribution="uniform"), trainable=trainable, name="lora_A")
self.B = tf.keras.layers.Dense(units=original_layer.output.shape[-1], kernel_initializer="zeros", trainable=trainable, name="lora_B")

def call(self, inputs):
original_outputs = self.original_layer(inputs)
lora_outputs = self.B(self.A(inputs))
return original_outputs + lora_outputs

lora_model._tracker.unlock()
lora_model.layers[-2] = LoraLayer(lora_model.layers[-2], rank=8, trainable=False)
lora_model._tracker.lock()
lora_model.summary()


My intention with this code is to replace the penultimate layer with a LoRA of itself. However, as per the model's summary, this is not happening:

The last layer is still a dense layer, and the number of trainable parameters does not change at all.

Is there something I am doing wrong in my implementation? Perhaps it has something to do with what is explained in this issue, and it is fundamentally impossible to do in Keras, or is there some intermediate step I am missing for getting this to work?

Lastly, I have also tried this solution, but the code is not working for me (it fails saying that a Node object has no attribute outbound_layer).

Continue reading...