How to fix some layers for transfer learning? #1706
-
|
Hi, Suppose I need do transfer learning based on pretrained ResNet50, and I want to fix the first layer, and allow other layers to update. How to do this in flax? I noticed #1176, they suggest to use https://flax.readthedocs.io/en/latest/flax.optim.html#flax.optim.MultiOptimizer, but there are two concerns, one is that it seems flax.optim tends to be replace by optax and optax do not have similar API. Second, In my project I already using optax, so I do not prefer change the code just for fix a layer, it seems cumbersome, I think there should be easy way to do it. The preferred way I can imagine is something like below:
|
Beta Was this translation helpful? Give feedback.
Replies: 2 comments 9 replies
-
|
Hi @wztdream, note that optax is now the recommended optimizer API (see here). You can freeze a subset of your params by using optax.multi_transform. One way of achieving this is to create a mask of your params, that assigns one label to trainable parameters and another label to frozen parameters. Consider this simple parameter tree: We want to freeze the params of the layers whose names start with And then: def zero_grads():
# from https://github.com/deepmind/optax/issues/159#issuecomment-896459491
def init_fn(_):
return ()
def update_fn(updates, state, params=None):
return jax.tree_map(jnp.zeros_like, updates), ()
return optax.GradientTransformation(init_fn, update_fn)
tx = optax.multi_transform({'adam': optax.adam(0.1), 'zero': zero_grads()},
create_mask(params, lambda s: s.startswith('frozen')))This simply means that all parameters with the label The only thing you have to do for this to work is give the modules you want to freeze a name that starts with I recently did something similar, so I created a Colab with a full example: https://colab.research.google.com/drive/1g_pt2Rc3bv6H6qchvGHD-BpgF-Pt4vrC#scrollTo=WWHlukuvIpXb Hope that helps! |
Beta Was this translation helpful? Give feedback.
-
|
thank you, it should be the recommended way to fix layers in flax. I just noticed there is a set_to_zero API in optax that did the same thing as the And also we can directly pass prefix of pytree as mask labels, which is easy to use, it will be something like: tx = optax.multi_transform({'adam': optax.adam(0.1), 'zero': zero_grads()},
{"params":{"frozen1":"zero", "trainable2":"adam", "trainable3":"adam", "trainable4":"adam"}})One problem in this case we need to from flax.core.frozen_dict import unfreeze
state = train_state.TrainState.create(apply_fn=model.apply,
params=unfreeze(params),
tx=tx) |
Beta Was this translation helpful? Give feedback.
-
|
Nice! You don't have to unfreeze the params if you freeze the mask: tx = optax.multi_transform({'adam': optax.adam(0.1), 'zero': zero_grads()},
frozen_dict.freeze({"params":{"frozen1":"zero", "trainable2":"adam", "trainable3":"adam", "trainable4":"adam"}}))
state = train_state.TrainState.create(apply_fn=model.apply,
params=params,
tx=tx) |
Beta Was this translation helpful? Give feedback.
-
|
See #1453 for a small recipe how to create the second parameter for |
Beta Was this translation helpful? Give feedback.
-
|
Thank you for the phenomenal answer. This has been an erudite discussion. I am using flax with Google's brax for reinforcement learning, and I am encountering a weird issue. Specifically, brax does the following: init_params = ppo_losses.PPONetworkParams(
policy=ppo_network.policy_network.init(key_policy),
value=ppo_network.value_network.init(key_value))
training_state = TrainingState( # pytype: disable=wrong-arg-types # jax-ndarray
optimizer_state=optimizer.init(init_params), # pytype: disable=wrong-arg-types # numpy-scalars
params=init_params,
normalizer_params=running_statistics.init_state(
specs.Array(env_state.obs.shape[-1:], jnp.dtype('float32'))),
env_steps=0)The line optimizer_state=optimizer.init(init_params), # pytype: disable=wrong-arg-types # numpy-scalarsworks great when using the PPONetworkParams dataclass. However, when I defined the masked optimizer as above, e.g. optimizer = optax.multi_transform(
{
"adam": optax.adam(3.0e-4),
"zero": zero_grads()
}, {
"Encoder_0": "zero",
"BCModel_0": "zero",
"ResidualPolicy_0": "adam"
})and then call Both optimizer variants are of type The full stack trace is here: |
Beta Was this translation helpful? Give feedback.
-
|
Hey @matthias-wright, Thanks a lot for the google colab - it's great! Do you know by any chance if this can save a significant amount of memory? In your example if the first layer is frozen it never needs to compute the gradients. E.g. ideally during the forward pass the activations should never be saved so that we can save some memomry. Is this the case here? Image you want to fine-tune a large BERT model and freeze the first 12 layers, but only train the final layer. In this case, it would be very important that no activations for the first 12 layers are computed to save memory. Is this the case here? |
Beta Was this translation helpful? Give feedback.
-
|
Hi @patrickvonplaten, I am happy to hear that the Colab was useful! Regarding your question, I am not entirely sure but I think it does. According to the docs of the gradient transformation However, I have never made any measurements to test this. Maybe you can test this by disabling the preallocation behavior of Jax and checking whether or not freezing some layers reduces the memory footprint? |
Beta Was this translation helpful? Give feedback.
-
|
One thing I tried in your Colab is to freeze all layers except the last one, then run a single training step in a Colab cell with This will show that freezing the layers is indeed faster! |
Beta Was this translation helpful? Give feedback.
-
|
I too had a play around with the Colab - it's brilliant! Taking inspiration from @marcvanzee, I compared three different approaches:
The Colab can be found here: https://colab.research.google.com/drive/1K-5bz6R6kt9GAvaUHvzYvvA-IOAO2PhL?usp=sharing Timing the jit'd train steps, freezing using Interesting to see that the train step time for |
Beta Was this translation helpful? Give feedback.
-
|
I think the reason why the stop_gradient solution isn't faster is because the conv kernels are so small that generating the zeros for the param gradient has comparable cost to just computing the convolution. If you make the number of features larger stop_gradient should be quite a bit faster than freezing no layers but still a bit slower than using optax.multi_transform |
Beta Was this translation helpful? Give feedback.
-
|
I modified the notebook to benchmark four different optimizers, training a 'tiny' seq2seq speech model comprising of ~50k parameters over 1000 training steps on TPU. Of these 50k parameters, the first ~17k are frozen. The training data is created synthetically with input/output features representative of real audio data:
The four optimizers trailed are:
The performance results for the three optimizers that freeze model parameters are very comparable:
Notebook: |
Beta Was this translation helpful? Give feedback.
Hi @wztdream,
note that optax is now the recommended optimizer API (see here).
You can freeze a subset of your params by using optax.multi_transform.
One way of achieving this is to create a mask of your params, that assigns one label to trainable parameters and another label to frozen parameters.
Consider this simple parameter tree:
We want to freeze the params of the layers whose names start with
frozen, so we assign them the labelzero(the name is arbitrary). The other parameters will be assigned the nameadam(also arbitrary):