How to do transfer learning with flaxmodels?

[riven314]

I am sorry to write another issue again because I am trying to apply transfer learning on ResNet using your flaxmodels.
However, I got stuck on how to get the backbone and add a new head on top of the ResNet.
Do you have any sample code/ guidance for me as a reference?

[riven314]

I saw example of transfer learning from flaxvision: https://github.com/rolandgvc/flaxvision/blob/master/examples/transfer_learning.ipynb
However, I feel like the script is not inheriting the pre-trained model weight in the training:

def get_initial_params(key):
  init_shape = jnp.ones((1, 224, 224, 3), jnp.float32)
  initial_params = MyModel().init(key, init_shape)['params']
  return initial_params

[matthias-wright]

Hi @riven314, check out the example I posted here.

[suniachar]

I tried to execute : https://github.com/rolandgvc/flaxvision/blob/master/examples/transfer_learning.ipynb
for my dataset, but I am getting initial accuracy as 14%, which is very low. What may be the reason?

[matthias-wright]

Hi @suniachar, you should ask the author of the flaxvision package. Best wishes