Transfer learning for Flax ResNet

[riven314]

I found a pretrained Flax ResNet in this repo: https://github.com/matthias-wright/flaxmodels
so I load in the pretrained ResNet with the following:

!pip install --upgrade git+https://github.com/matthias-wright/flaxmodels.git

import numpy as np
import jax
import jax.numpy as jnp
from flaxmodels import ResNet18, ResNet34

# num_classes must be 1000 if pretrained
resnet = ResNet18(output='softmax', pretrained='imagenet')
key = jax.random.PRNGKey(0)
# (batch, width, height, channel)
im = jnp.array(np.random.rand(1, 480, 360, 3))
params = resnet.init(key, im)

After that, I would like to do transfer learning (i.e. freeze the body weight and replace a fresh head) on the pretrained ResNet.
However, I had difficulty finding relevant reference code in Flax. Could you give me some guidance/ sample code on how I could achieve that in Flax?

[matthias-wright]

Hi @riven314, check out my answer in #1706.
There is also an accompanying Colab, where I show how to freeze a subset of layers for transfer learning.
Hope that helps!