Is it possible to adapt the shape of trained weights after extending the classification head?

Question

I have an image classification tensorflow model that's based on mobilenetV2, but with my own classification head. I've trained it and now it works fine so far.

Every now and then i need to add further classes though. Since training is quite time consuming with my hardware, i'd like to keep as much of the already trained weights. At the moment i just replace the classification head with a new one. So the weights of all layers are preserved, except those of the classification head.

My question now is:
Is it possible to also keep the weights of the classification head and just extend them by an empty weight for every new class?

I'v already tried to add the former weights to the new classification head, but that throws an error saying: "Layer weight shape ... not compatible with provided weight shape ..." Is there a way to extend the weight shape so that it matches with the shape of the new classification head?

Answer 1

I came to a solution, which seems to do the trick. This function (typescript / tensorflow.js) extends the classification head by one class and adapts the trained weights to the new shape:

const addClass = (model: tf.LayersModel): tf.LayersModel => {
  const formerHead = model.layers[model.layers.length - 1];
  const formerHeadWeights = formerHead.getWeights();
  const newOutput = tf.layers.dense({
    units: (formerHead as any).units + 1,
    activation: 'softmax',
    weights: [
      tf.tensor((formerHeadWeights[0].arraySync() as number[][]).map((row) => [...row, 0])),
      tf.tensor([...(formerHeadWeights[1].arraySync() as number[]), 0])
    ]
  })
    .apply(model.layers[model.layers.length - 2].output) as tf.SymbolicTensor;

  return tf.model({inputs: model.inputs, outputs: newOutput});
};