How to freeze layer or detach a layer in training stage.

[haoweiz23]

I have read the flax and jax docs. I have not found api can be used to freeze or detach some layers in training stage.
In pytorch, this api is attribute 'require_grad'. In keras, this api is attribute 'trainable'.
My mission is freeze some layers in my model at beginning and train the whole model after a while. So I need to freeze some layers and avoid parameters update.

Does there any examples or function attributes can be used to deal with this problem?

[avital]

Hi @Zhu-haow, BTW we typically use GitHub Discussions for questions.

Take a look at the focus argument to OptimizerDef.create. You can easily define Traversals using optim.ModelParamTraversal, for example:

      kernels = optim.ModelParamTraversal(lambda path, _: 'kernel' in path)
      biases = optim.ModelParamTraversal(lambda path, _: 'bias' in path)

(We intentionally don't make the property stateful on the modules because you may want to optimize different subsets at different opportunities)

cc @marcvanzee for an example of a few bits of missing documentation (also note that the MultiOptimizer has a docstring for the constructor, but we don't show that on RTD), and perhaps a useful HOWTO

[andsteing]

Note that by now flax.optim is deprecated; see #1453 for an answer in how to do this with Optax.