# Re-using layers in Flux.jl: how to train a multi-layer model sharing a common LSTM layer and separate dense layers?

**URL:** <https://discourse.julialang.org/t/re-using-layers-in-flux-jl-how-to-train-a-multi-layer-model-sharing-a-common-lstm-layer-and-separate-dense-layers/82705>\
**Category:** Machine Learning\
**Created:** [June 13, 2022, 9:10pm UTC](https://discourse.julialang.org/t/re-using-layers-in-flux-jl-how-to-train-a-multi-layer-model-sharing-a-common-lstm-layer-and-separate-dense-layers/82705 "2022-06-13T21:10:43Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![rio](https://avatars.discourse-cdn.com/v4/letter/r/ea666f/32.png) [@rio](https://discourse.julialang.org/u/rio)\
**Post date:** [June 13, 2022, 9:10pm UTC](https://discourse.julialang.org/t/re-using-layers-in-flux-jl-how-to-train-a-multi-layer-model-sharing-a-common-lstm-layer-and-separate-dense-layers/82705/1 "2022-06-13T21:10:44Z")

</div>

I would like to train a model that consists of two “submodels”, say model1 and model2. The two models share a common LSTM layer and are followed by two separate dense layers. The final output is the average of the two models.

My (failed) attempt:

```julia
using Flux
using Flux.Data: DataLoader
using Flux.Optimise: Optimiser 

nfeatures = 11
seq_length = 28
nsamples = 1000
opt = ADAM(0.002)
device = cpu

# Generete toy samples
X = rand(Float32, nfeatures, seq_length, nsamples)
Y = sum(X[1, end-1:end, :], dims=1)

train_loader = DataLoader((X, Y), batchsize=64, shuffle=true)

function apply_model(model, x)

    model1 = Chain(model[:mycommomlayer], model[:mydense1])
    model2 = Chain(model[:mycommomlayer], model[:mydense2])

    Flux.reset!(model1)
    y_1 = last(map(model1, [view(x, :, t, :) for t in 1:seq_length]))

    Flux.reset!(model2)
    y_2 = last(map(model2, [view(x, :, t, :) for t in 1:seq_length]))

    return (y_1 + y_2)/2
end

model = Chain(mycommomlayer = LSTM(nfeatures, 32),
    mydense1 = Chain(Dropout(0.1), Dense(32, 1)),
    mydense2 = Chain(Dropout(0.1), Dense(32, 1)))

function loss(ŷ, y)
    l = Flux.mse(y, ŷ)
    println("loss is $l")
    return l
end

ps = Flux.params(model)

for epoch in 1:5
    @info(" start epoch $epoch ")
    for (x, y) in train_loader
        x, y = x |> device, y |> device
        gs = Flux.gradient(ps) do
            ŷ = apply_model(model, x)
            loss(ŷ, y)
        end
        Flux.Optimise.update!(opt, ps, gs)
    end
end

```

It crashes with

> [ Info: start epoch 1  
> loss is 1.6848714  
> ERROR: MethodError: no method matching +(::Base.RefValue{Any}, ::NamedTuple{(:cell, :state), Tuple{Nothing, Tuple{Matrix{Float32}, Matrix{Float32}}}})  
> Closest candidates are:  
> +(::Any, ::Any, ::Any, ::Any…) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/base/operators.jl:655  
> +(::ChainRulesCore.AbstractThunk, ::Any) at ~/.julia/packages/ChainRulesCore/ksvfu/src/tangent\_arithmetic.jl:122  
> +(::ChainRulesCore.Tangent{P}, ::P) where P at ~/.julia/packages/ChainRulesCore/ksvfu/src/tangent\_arithmetic.jl:146  
> …  
> Stacktrace:  
> [1] accum(x::Base.RefValue{Any}, y::NamedTuple{(:cell, :state), Tuple{Nothing, Tuple{Matrix{Float32}, Matrix{Float32}}}})  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:17  
> [2] macro expansion  
> @ ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:57 [inlined]  
> [3] accum\_param(cx::Zygote.Context, x::Flux.Recur{Flux.LSTMCell{Matrix{Float32}, Vector{Float32}, Tuple{Matrix{Float32}, Matrix{Float32}}}, Tuple{Matrix{Float32}, Matrix{Float32}}}, Δ::NamedTuple{(:cell, :state), Tuple{Nothing, Tuple{Matrix{Float32}, Matrix{Float32}}}})  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:53  
> [4] back  
> @ ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:228 [inlined]  
> [5] #1765#back  
> @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]  
> [6] Pullback  
> @ ./namedtuple.jl:127 [inlined]  
> [7] #213  
> @ ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:203 [inlined]  
> [8] (::Zygote.var"#1754#back#215"{Zygote.var"#213#214"{Tuple{Tuple{Nothing}, Tuple{Nothing}}, typeof(∂(getindex))}})(Δ::NamedTuple{(:cell, :state), Tuple{Nothing, Tuple{Matrix{Float32}, Matrix{Float32}}}})  
> @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67  
> [9] Pullback  
> @ ~/.julia/packages/MacroTools/PP9IQ/src/examples/forward.jl:18 [inlined]  
> [10] (::Zygote.var"#213#214"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(#getindex#149))})(Δ::NamedTuple{(:cell, :state), Tuple{Nothing, Tuple{Matrix{Float32}, Matrix{Float32}}}})  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:203  
> [11] #1754#back  
> @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]  
> [12] Pullback  
> @ ~/.julia/packages/MacroTools/PP9IQ/src/examples/forward.jl:18 [inlined]  
> [13] Pullback  
> @ ~/.julia/packages/Zygote/cCyLF/src/tools/builtins.jl:15 [inlined]  
> [14] (::typeof(∂(literal\_getindex)))(Δ::NamedTuple{(:cell, :state), Tuple{Nothing, Tuple{Matrix{Float32}, Matrix{Float32}}}})  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0  
> [15] Pullback  
> @ ~/Library/CloudStorage/Box-Box/MZ/proj/streamflow\_ungauged/finetunning\_physic/toy.jl:21 [inlined]  
> [16] (::typeof(∂(apply\_model)))(Δ::Matrix{Float32})  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0  
> [17] Pullback  
> @ ~/Library/CloudStorage/Box-Box/MZ/proj/streamflow\_ungauged/finetunning\_physic/toy.jl:49 [inlined]  
> [18] (::typeof(∂(λ)))(Δ::Float32)  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0  
> [19] (::Zygote.var"#94#95"{Zygote.Params, typeof(∂(λ)), Zygote.Context})(Δ::Float32)  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface.jl:357  
> [20] gradient(f::Function, args::Zygote.Params)  
> @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface.jl:76  
> [21] top-level scope  
> @ ~/Library/CloudStorage/Box-Box/MZ/proj/streamflow\_ungauged/finetunning\_physic/toy.jl:48

Please, how could I define the two models that share the LSTM layer?  
Many thanks in advance!

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [June 15, 2022, 2:16am UTC](https://discourse.julialang.org/t/re-using-layers-in-flux-jl-how-to-train-a-multi-layer-model-sharing-a-common-lstm-layer-and-separate-dense-layers/82705/2 "2022-06-15T02:16:19Z")

</div>

Right now, you are also mapping the dense layer over every element in the sequence. I presume that is not intended? If not, then the loss function should first apply `mycommonlayer` over x twice to get `y_1` and `y_1`, _then_ call `mydense1` and `mydense2` on those.

You can pass `mycommonlayer`, `mydense1` and `mydense2` into your loss function (you can toss them into a (named)tuple if you still want `model` as a single arg). `mycommonlayer` will be automatically shared.

Secondly, I would recommend pre-arranging your input into a vector of views. Doing it in the loss where the AD can see it will be prohibitively slow.

Putting all this together:

```julia
function apply_model(rnn, dense1, dense2, x)
    Zygote.ignore(() -> Flux.reset!(rnn)) # we don't need to differentiate this
    y_1 = [rnn(step) for step in x][end]

    Zygote.ignore(() -> Flux.reset!(rnn)) # we don't need to differentiate this
    y_2 = [rnn(step) for step in x][end]

    return (mydense1(y_1) .+ mydense2(y_2)) ./ 2
end

```

Alternatively, you could package up the `(dense + dense) / 2` part into a single layer:

```julia
addscale(y_1, y_2) = (y_1 .+ y_2) ./ 2
post_rnn = Parallel(addscale, mydense1, mydense2)

```