# How do I retain model parameters while freezing layers with Lux.jl

**URL:** <https://discourse.julialang.org/t/how-do-i-retain-model-parameters-while-freezing-layers-with-lux-jl/105558>\
**Category:** Machine Learning\
**Tags:** lux\
**Created:** [October 30, 2023, 3:28am UTC](https://discourse.julialang.org/t/how-do-i-retain-model-parameters-while-freezing-layers-with-lux-jl/105558 "2023-10-30T03:28:19Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![liuyxpp](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/liuyxpp/32/9870_2.png) [@liuyxpp](https://discourse.julialang.org/u/liuyxpp)\
**Post date:** [October 30, 2023, 3:28am UTC](https://discourse.julialang.org/t/how-do-i-retain-model-parameters-while-freezing-layers-with-lux-jl/105558/1 "2023-10-30T03:28:19Z")

</div>

I am very interested in the example given in the doc [Freezing Model Parameters | LuxDL Docs](https://lux.csail.mit.edu/v0.5.9/manual/freezing_model_parameters#freezing-part-of-a-chain). However, in that example, the parameters are obtained from

```julia
ps, st = Lux.setup(rng, model_frozen)

```

I find there is no convenient way to retain the parameters of the old model. Say I have a pre-trained model, which is a `Chain` of nested `Dense` and `GRU` layers. The trained parameters and states are `ps_trained` and `st_trained`. How can I use the method provided in the doc to use `ps_trained` and `st_trained`?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 30, 2023, 10:43pm UTC](https://discourse.julialang.org/t/how-do-i-retain-model-parameters-while-freezing-layers-with-lux-jl/105558/2 "2023-10-30T22:43:53Z")

</div>

Check the sections just before that. It should how to use [Experimental Features | LuxDL Docs](https://lux.csail.mit.edu/v0.5.9/api/Lux/contrib#Lux.Experimental.@layer_map) to accomplish exactly what you asked for

---

<div class="post-metadata">

**Author:** ![liuyxpp](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/liuyxpp/32/9870_2.png) [@liuyxpp](https://discourse.julialang.org/u/liuyxpp)\
**Post date:** [October 31, 2023, 1:37am UTC](https://discourse.julialang.org/t/how-do-i-retain-model-parameters-while-freezing-layers-with-lux-jl/105558/3 "2023-10-31T01:37:02Z")

</div>

Yeah, I have tried `@layer_map` before. However, it seems not working properly with `Recurrence(cell = GRUCell(32 => 32))` layer. The error message is

```julia
ERROR: AssertionError: fieldnames(typeof(st_c)) == fieldnames(typeof(ps_c))
Stacktrace:
 [1] layer_map(f::typeof(freeze_except_last_layer), l::Recurrence{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, name::String)
   @ Lux.Experimental ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:96
 [2] layer_map(f::typeof(freeze_except_last_layer), l::@NamedTuple{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, name::String)
   @ Lux.Experimental ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:108
 [3] layer_map(f::typeof(freeze_except_last_layer), l::Lux.Chain{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, name::String)
   @ Lux.Experimental ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:108
 [4] layer_map(f::typeof(freeze_except_last_layer), l::SkipConnection{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, name::String)
   @ Lux.Experimental ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:108
 [5] layer_map(f::typeof(freeze_except_last_layer), l::@NamedTuple{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, name::String)
   @ Lux.Experimental ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:108
 [6] layer_map(f::typeof(freeze_except_last_layer), l::Lux.Chain{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, name::String)
   @ Lux.Experimental ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:108
 [7] top-level scope
   @ ~/.julia/packages/Lux/Al3Ab/src/contrib/map.jl:38
Some type information was truncated. Use `show(err)` to see complete types.

```

I believe there is a bug in `@layer_map`.

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 31, 2023, 1:52pm UTC](https://discourse.julialang.org/t/how-do-i-retain-model-parameters-while-freezing-layers-with-lux-jl/105558/4 "2023-10-31T13:52:28Z")

</div>

Can you open an issue with a reproducer?