# Modify existing Lux models
**URL:** https://discourse.julialang.org/t/modify-existing-lux-models/105199
**Category:** Machine Learning
**Tags:** question, neural-network, lux
**Created:** [October 19, 2023, 7:36pm UTC](https://discourse.julialang.org/t/modify-existing-lux-models/105199 "2023-10-19T19:36:53Z")
**Posts on this page:** 5
**Page:** 1
### Author:  [@Andreu](https://discourse.julialang.org/u/Andreu)
#### Post date: [October 19, 2023, 7:36pm UTC](https://discourse.julialang.org/t/modify-existing-lux-models/105199/1 "2023-10-19T19:36:54Z")
Is there a systematic approach to obtain a modified version of an existing Lux model?
The [documentation](https://lux.csail.mit.edu/dev/manual/interface#layer-interface) of the Lux layer interface is clear: “once constructed a model architecture cannot change.” However, we may want to replace or modify parts of an existing, possibly pretrained Lux model.
An example of replacement could be using the backbone of a vision model as a feature extractor while replacing the classifier layers to address a different task. An example of modification could be [this proposed solution](https://stackoverflow.com/a/54777347/1306284) to adapt a vision model trained on RGB images to work on grayscale images (adapting both the architecture and the pretrained parameters of the input layer).
I see that I can manually create a new Lux model with the required changes. In the minimal working example below, `model_prime` is like `model` but changing the second layer (both architecture and parameters):
```julia
using Lux
using Random
rng = Random.default_rng()
Random.seed!(rng, 0)
model = Lux.Chain(
Lux.Dense(4, 3),
Lux.Dense(3, 1)
)
ps, st = Lux.setup(rng, model)
x = randn(rng, Float32, 4, 1);
@show size(model(x, ps, st)[1])
model_prime = Lux.Chain(
Lux.Dense(4, 3),
Lux.Dense(3, 6)
)
ps_prime = merge(
ps,
(layer_2 = (weight = randn(rng, Float32, 6, 3), bias = randn(rng, Float32, 6, 1)), )
)
@show size(model_prime(x, ps_prime, st)[1])
```
However, this seems inconvenient and even error-prone as soon as the model becomes more complex. To achieve the goal in a more systematic manner, I could think of copying the model and modifying only the layers of interest. However, in my minimal working example the following code
```julia
model_copy = deepcopy(model)
model_copy.layers.layer_2 = Lux.Dense(3, 6)
```
throws an `ERROR: setfield!: immutable struct of type NamedTuple cannot be changed`.
I would appreciate advice. Should I investigate how Lux models are defined under the hood to find a mechanism? Is there an obvious way to achieve this goal?
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [October 19, 2023, 8:08pm UTC](https://discourse.julialang.org/t/modify-existing-lux-models/105199/2 "2023-10-19T20:08:03Z")
Use `Setfield.jl`
`@set! model.layers.layer_2 = Dense(3 => 6)`
Same thing can be done for the parameters.
---
### Author:  [@liuyxpp](https://discourse.julialang.org/u/liuyxpp)
#### Post date: [October 20, 2023, 2:19am UTC](https://discourse.julialang.org/t/modify-existing-lux-models/105199/3 "2023-10-20T02:19:55Z")
Hi @avikpal I think my question is relevant to this topic so I’d like to ask it here to see whether what I’ve done is reasonable. My aim is to implement a RNN with `return_sequence=True` and its input and output should be preprocessed and post-processed by Dense layers, like the simple model below:
```julia
function GRU_Seq2Seq(nfeatures; gru_out_dim=GRU_OUT, dense_act=DENSE_ACT)
dense_in1 = Lux.Dense(nfeatures=>gru_out_dim*2, dense_act)
drop1 = Lux.Dropout(0.5)
dense_in2 = Lux.Dense(gru_out_dim*2=>gru_out_dim, dense_act)
gru = Lux.Recurrence(Lux.GRUCell(gru_out_dim=>gru_out_dim); return_sequence=true)
stack_layer = StackLayer(2)
dense_out1 = Lux.Dense(gru_out_dim=>16, dense_act)
drop2 = Lux.Dropout(0.2)
dense_out2 = Lux.Dense(16=>1, identity)
return Lux.Chain(dense_in1, drop1, dense_in2, gru, stack_layer, dense_out1, drop2, dense_out2)
end
```
where the `StackLayer` is defined to make the RNN output suitable for Dense layers:
```julia
struct StackLayer <: Lux.AbstractExplicitLayer
dims::Int
end
@inline function (s::StackLayer)(x::AbstractArray, ps, st::NamedTuple)
return stack(x, dims=s.dims), st
end
```
I’d like to know is there any idiomatic way to define such Seq2Seq RNN model without hacking my own `StackLayer`. Do I miss something from the Lux.jl doc? I feel it is a natural thing to do a post processing by Dense layers.
Moreover, to test the performance of my seq2seq model, I also implement a seq2last model which only return the last value like following:
```julia
function GRU_Seq2Last(nfeatures; gru_out_dim=GRU_OUT, dense_act=DENSE_ACT)
dense_in1 = Lux.Dense(nfeatures=>gru_out_dim*2, dense_act)
drop1 = Lux.Dropout(0.5)
dense_in2 = Lux.Dense(gru_out_dim*2=>gru_out_dim, dense_act)
gru = Lux.Recurrence(Lux.GRUCell(gru_out_dim=>gru_out_dim))
dense_out1 = Lux.Dense(gru_out_dim=>16, dense_act)
drop2 = Lux.Dropout(0.2)
dense_out2 = Lux.Dense(16=>1, identity)
return Lux.Chain(dense_in1, drop1, dense_in2, gru, dense_out1, drop2, dense_out2)
end
```
Now that the `GRU_Seq2Last` is trained, I’d like use it to evaluate the test data but return a full sequence instead of the last value as in the training phase. Currently, I use a really ugly hack to transform `GRU_Seq2Last` parameters and states to be usable by `GRU_Seq2Seq` and use this model to predict. The codes are:
```julia
function seq2last_to_seq2seq(ps, st)
ps_out = (layer_1=ps.layer_1, layer_2=ps.layer_2,
layer_3=ps.layer_3, layer_4=ps.layer_4,
layer_5=NamedTuple(),
layer_6=ps.layer_5, layer_7=ps.layer_6,
layer_8=ps.layer_7)
st_out = (layer_1=st.layer_1, layer_2=st.layer_2,
layer_3=st.layer_3, layer_4=st.layer_4,
layer_5=NamedTuple(),
layer_6=st.layer_5, layer_7=st.layer_6,
layer_8=st.layer_7)
return ps_out, st_out
end
# ps_seq2last, st_seq2last is trained by GRU_Seq2Last
ps_seq2seq, st_seq2seq = seq2last_to_seq2seq(ps_seq2last, ps_seq2last)
yp_seq2seq = GRU_Seq2Seq(40)(X, ps_seq2seq, st_seq2seq)
```
The Setfield.jl seems can’t help here since I want to insert a layer (the `StackLayer`) instead of replace it.
---
### Author:  [@Andreu](https://discourse.julialang.org/u/Andreu)
#### Post date: [October 23, 2023, 11:44am UTC](https://discourse.julialang.org/t/modify-existing-lux-models/105199/4 "2023-10-23T11:44:31Z")
Thank you, @avikpal. This solves my question. Let me add though that, after reading the README of Setfield.jl, I have tried your proposed solution but only with the successor package Accessors.jl, not with Setfield.jl.
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [October 23, 2023, 3:30pm UTC](https://discourse.julialang.org/t/modify-existing-lux-models/105199/5 "2023-10-23T15:30:34Z")
Great! Both should work, but I run tests only with Setfield so I recommended that.