Реализация оболочки LoRA для Conv2D в Tensorflow ⇐ Python

[Anonymous]

Недавно я пытался портировать метод LoRA для CNN, первоначально реализованный в PyTorch, в TensorFlow, но моя реализация не сработала, и мне нужна помощь, чтобы заставить ее работать.
Мой код можно найти здесь (он относительно обширен, поэтому я изложил его вкратце). Короче говоря, метод LoRA, который я реализовал, работает аналогично тому, который уже реализован в слоях Conv в TF, но я хотел разрешить включение или отключение LoRA. Я реализовал слой LoRAConv2D, который окружал обычный слой Conv2D и добавил пару матриц A и B в качестве весов, которые будут использоваться в LoRA:

Код: Выделить всё

def build():
...
self.lora_A: tf.Variable = self.add_weight(
name="LoRA_matA",
shape=lora_A_shape,
initializer=keras.initializers.HeUniform(),
trainable=self.lora_enabled,
)

self.lora_B: tf.Variable = self.add_weight(
name="LoRA_matB",
shape=lora_B_shape,
initializer=keras.initializers.Zeros(),
trainable=self.lora_enabled,
)

Затем эти матрицы будут использоваться при прямом выполнении слоя, в зависимости от того, включено ли использование LoRA:

Код: Выделить всё

def call(self, inputs):
if self.merged or not self.lora_enabled:
return self.__original_layer(inputs)

delta_W = tf.matmul(self.lora_B, self.lora_A)
delta_W = tf.reshape(delta_W, self.kernel_shape)

outputs = tf.nn.conv2d(
inputs,
self.__original_layer.kernel + self.scaling * delta_W,
strides=self.__original_layer.strides,
padding=self.__original_layer.padding.upper(),
)

if self.__original_layer.use_bias:
outputs = tf.nn.bias_add(outputs, self.__original_layer.bias)

if self.__original_layer.activation is not None:
outputs = self.__original_layer.activation(outputs)

return outputs

Я также реализовал функции для установки соответствующего обучаемого атрибута на матрицах и в __original_layer.kernel, чтобы контролировать, что именно следует обучать в зависимости от того, используется LoRA или нет.
К сожалению, эта реализация работала не совсем. Все относительно нормально, но при выполнении тестового запуска с использованием образца модели (функция test_lora_conv2d в сути) получаю следующее предупреждение:

Код: Выделить всё

UserWarning: Gradients do not exist for variables ['lo_ra_conv2d/LoRA_matA', 'lo_ra_conv2d/LoRA_matB', 'lo_ra_conv2d_1/LoRA_matA', 'lo_ra_conv2d_1/LoRA_matB'] when minimizing the loss. If using `model.compile()`, did you forget to provide a `loss` argument?

Дальнейшая проверка показала мне, что когда я создаю и строю модель со слоями LoRAConv2D, у которых для self.lora_enabled установлено значение False, матрицы не получают градиенты при вызове model.fit(), даже если я включил LoRA и скомпилировал модель до вызова функции. Разве установка переменных как обучаемых и последующая перекомпиляция модели не должна позволять им получать градиенты и обрабатываться оптимизатором? Я здесь что-то делаю неправильно?

Подробнее здесь: https://stackoverflow.com/questions/798 ... tensorflow