Why do the parameters of a network is changing when i make gradient =0?

[maobenz]

I use the stop_gradient() to make some paramters of a network be 0, and i print it out and it proves they are all zero.

However, after some iterations, the value of parameters is changing.

Code:

Class A:
    Class B:
    Class C:


   def _call_:
        b = B(input)
        c = C(b)
        c = stop_gradient(c)

print(grad)
## grad of model c is all zero
print(state.optimizer.target['model']['c'])
## after one iterations
the value of the paramter is changing

Why?? is the reason that the value of parameters of model b is changing??
How do i really stop model c learn??
Thanks

[YouJiacheng]

What optimizer did you use?

[YouJiacheng]

I can't clearly figure out the problem, but at least you can manually copy the parameters of model C from the old state.
flax.optim is deprecated and I don't know how to work with it. If you use optax I can offer some help.
see https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md

[maobenz]

Thanks. i will try the optax. Actually maybe i can't manually copy the paramters because when i get the parameters from the optimizer, it is the frozen dict type and it can't be changed.

[YouJiacheng]

There are unfrozen and copy(with replace or add) APIs for frozen dict.

[maobenz]

Thanks!
I think unfrozen and copy api just create new dictionary but not change the orignal dict. And the orginal dict is also frozen.
just like
dic = dic.copy(params)
it doesn't work

[YouJiacheng]

Yes it creates a new dictionary, but you don't need to modify the original dict! You only need to modify the optimizer state!