Train only bias in a pre-trained model from Keras-Tensorflow

Question

Sounds strange, I know! but: Is possible train only the bias? I have pretrained models but applying low rank to the weights, obvously, the accurracy of the NN downs... Is there anyway to say Keras-TensoFlow compiler that only trains the bias? Certainly, I don't know if this has really sense (I suposse that is a silly idea...), but I want test if the accuracy increases or not..

Answer 1

You can manually select the variables to update when applying gradients like so:

def get_grad(model, x, y):
    with tf.GradientTape() as tape:
        loss = compute_loss(model, x, y, training=True)
        to_update = [i for ix, i in enumerate(model.trainable_variables) if ix in (1, 3, 5, 7)]
    return loss, tape.gradient(loss, to_update)

It returns the variables 1, 3, 5, 7 which are the biases. And yes, it does work:

Epoch 10 Loss: 1.013 Acc: 33.33%
Epoch 11 Loss: 1.006 Acc: 34.00%
Epoch 12 Loss: 0.999 Acc: 34.00%
Epoch 13 Loss: 0.993 Acc: 36.00%
Epoch 14 Loss: 0.987 Acc: 39.33%
Epoch 15 Loss: 0.982 Acc: 48.67%
Epoch 16 Loss: 0.979 Acc: 53.33%
Epoch 17 Loss: 0.975 Acc: 56.00%
Epoch 18 Loss: 0.972 Acc: 59.33%
Epoch 19 Loss: 0.969 Acc: 60.67%
Epoch 20 Loss: 0.967 Acc: 61.33%
Epoch 21 Loss: 0.966 Acc: 61.33%
Epoch 22 Loss: 0.962 Acc: 61.33%
Epoch 23 Loss: 0.961 Acc: 62.67%
Epoch 24 Loss: 0.961 Acc: 62.00%
Epoch 25 Loss: 0.959 Acc: 62.67%

Full code:

import tensorflow as tf
from tensorflow.keras.layers import Dense
from sklearn.datasets import load_iris
import numpy as np

X, y = load_iris(return_X_y=True)
X = X.astype(np.float32)

train = tf.data.Dataset.from_tensor_slices((X, y)).shuffle(25).batch(8)

model = tf.keras.Sequential([
    Dense(16, activation='relu'),
    Dense(32, activation='relu'),
    Dense(64, activation='relu'),
    Dense(3, activation='softmax')])

loss_object = tf.losses.SparseCategoricalCrossentropy(from_logits=False)


def compute_loss(model, x, y, training):
  out = model(x, training=training)
  loss = loss_object(y_true=y, y_pred=out)
  return loss


def get_grad(model, x, y):
    with tf.GradientTape() as tape:
        loss = compute_loss(model, x, y, training=True)
        to_update = [i for ix, i in enumerate(model.trainable_variables) if ix in (1, 3, 5, 7)]
    return loss, tape.gradient(loss, to_update)


optimizer = tf.optimizers.Adam()

verbose = "Epoch {:2d} Loss: {:.3f} Acc: {:.2%}"

model.build(input_shape=([None, 4]))
weights_before = model.layers[0].get_weights()

for epoch in range(1, 25 + 1):
    train_loss = tf.metrics.Mean()
    train_acc = tf.metrics.SparseCategoricalAccuracy()

    for x, y in train:
        loss_value, grads = get_grad(model, x, y)
        to_update = [i for ix, i in enumerate(model.trainable_variables) if ix in (1, 3, 5, 7)]
        optimizer.apply_gradients(zip(grads, to_update))
        train_loss.update_state(loss_value)
        train_acc.update_state(y, model(x, training=True))

    print(verbose.format(epoch,
                         train_loss.result(),
                         train_acc.result()))

weights_after = model.layers[0].get_weights()