What is the recommended way to do transfer learning in flax?

[wztdream]

There are already several discussion about transfer learning here, and here, but it seems not clear enough, may be due to their special use case.

Let's take a common use case as example, suppose we have ResNet50 pretrained on ImageNet(I mean the model using flax, for example ResNet50 in flax examples), and we want to use it as start point to train our own classification network, there are two operations we need to make:

1. remove some layers, for example remove the classification layer
2. freeze some layers, for example freeze first two layers, and make other layers trainable

Then we add our own classification layer on top of it to train on our dataset

How to do these in flax? It will be better to add an example for this common use case.

[andsteing]

Note that we also have a full transfer learning example (vision_transformer) mentioned in our examples RTD page:
https://flax.readthedocs.io/en/latest/examples.html

But still I think what you propose above could be a helpful addition, created #1567

[murphyk]

I don't see any transfer learning examples listed at https://flax.readthedocs.io/en/latest/examples.html.
I did find https://www.kaggle.com/code/yashvi/transfer-learning-using-jax-flax/ and https://www.machinelearningnuggets.com/transfer-learning-with-jax-flax/.
But a 'canonical' implementation from the flax team would be useful.

[cgarciae]

We are working on a Transfer Learning guide, for future users you can check it out here:

https://flax.readthedocs.io/en/latest/guides/transfer_learning.html

In the meanwhile you can check out the unreleased version here:

https://flax--2394.org.readthedocs.build/en/2394/guides/transfer_learning.html