# Neural Nets training with multiple Chains Lux.jl and CUDA.jl

**URL:** <https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167>\
**Category:** GPU\
**Created:** [October 3, 2022, 4:29pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167 "2022-10-03T16:29:44Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![LeoCott](https://avatars.discourse-cdn.com/v4/letter/l/94ad74/32.png) [@LeoCott](https://discourse.julialang.org/u/LeoCott)\
**Post date:** [October 3, 2022, 4:29pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/1 "2022-10-03T16:29:44Z")

</div>

Hi all,

I am a GPU beginner so I probably ask something obvious.

I defined a neural net with four Lux.jl chains:

```julia
input_ = 2
n = 20
chain1 = Lux.Chain(Dense(input_,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,1))
chain2 = Lux.Chain(Dense(input_,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,1))
chain3 = Lux.Chain(Dense(input_,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,1))
chain4 = Lux.Chain(Dense(input_,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,1))

```

And I would like to train it on a GPU. For that purpose, I ensure that the initial parameters are on the GPU as in the example [here](https://neuralpde.sciml.ai/stable/tutorials/gpu/)

```julia
ps1 = Lux.setup(Random.default_rng(), chain1)[1]
ps1 = ps1 |> Lux.ComponentArray |> gpu .|> Float64

ps2 = Lux.setup(Random.default_rng(), chain2)[1]
ps2 = ps2 |> Lux.ComponentArray |> gpu .|> Float64

ps3 = Lux.setup(Random.default_rng(), chain3)[1]
ps3 = ps3 |> Lux.ComponentArray |> gpu .|> Float64

ps4 = Lux.setup(Random.default_rng(), chain4)[1]
ps4 = ps4 |> Lux.ComponentArray |> gpu .|> Float64

```

And finally I use the symbolic discretization provided by NeuralPDE:

```julia
discretization = NeuralPDE.PhysicsInformedNN([chain1 , chain2, chain3, chain4],training_strategy, init_params = [ps1 ps2 ps3 ps4], param_estim=true, additional_loss=additional_loss)

```

When doing this, I have the following error _**AssertionError: length(init\_params) == length(depvars)**_, I also tried to replace init\_params = [ps1 ps2 ps3 ps4] by init\_params = [ps1, ps2, ps3, ps4] but had the error _ **CuArray only supports element types that are stored inline** _…

Any ideas ?

Thanks in advance!

---

<div class="post-metadata">

**Author:** ![de-souza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/de-souza/32/43417_2.png) [@de-souza](https://discourse.julialang.org/u/de-souza)\
**Post date:** [October 3, 2022, 6:43pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/2 "2022-10-03T18:43:10Z")

</div>

Hi, I faced a similar problem. Somebody else also created [an issue on GitHub](https://github.com/SciML/NeuralPDE.jl/issues/594).

Here is how I fixed it:

1. Define all the initial parameters as a single ComponentArray.

```julia
using ComponentArrays, CUDA, Lux, Random

# [...]

@parameters x y
@variables f1(..) f2(..) f3(..) f4(..)

# [...]

chain = [chain1 , chain2, chain3, chain4]
names = :f1, :f2, :f3, :f4 # same as the variables from the beginning

init_params = Lux.initialparameters.(Random.default_rng(),
                                     chain)
init_params = NamedTuple{names}(init_params)
init_params = ComponentArray(init_params)

```

> **Edit: Step 2 is no longer necessary as of ComponentArrays@v0.13.3.**
>
> 1. Redefine the conversion that happen at [NeuralPDE.jl/src/discretize.jl#L480](https://github.com/SciML/NeuralPDE.jl/blob/v5.3.0/src/discretize.jl#L480) (necessary for me as of NeuralPDE@v5.3.0 and ComponentArrays@v0.13.2).
> 
> ```julia
> using ComponentArrays: GPUComponentArray
> 
> function ComponentArray(nt::NamedTuple{(:depvar,),
> <:Tuple{GPUComponentArray{T}}}) where {T}
> depvar = cpu(nt.depvar)
> A = ComponentArray(; depvar)
> A = T.(gpu(A))
> return A
> end
> 
> ```

1. Move the initial parameters to the GPU.

```julia
init_params = Float64.(gpu(init_params))

```

Let me know if this works for you. I will try to submit a PR to NeuralPDE with the fix.

---

<div class="post-metadata">

**Author:** ![LeoCott](https://avatars.discourse-cdn.com/v4/letter/l/94ad74/32.png) [@LeoCott](https://discourse.julialang.org/u/LeoCott)\
**Post date:** [October 4, 2022, 2:59pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/3 "2022-10-04T14:59:48Z")

</div>

Hi, thanks a lot for your answer @de-souza !

Unfortunately, I still have the same issue… Is there a simpler way to define the initial parameters as an array of gpu initialized values as mentioned in the issue on GitHub ?

---

<div class="post-metadata">

**Author:** ![de-souza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/de-souza/32/43417_2.png) [@de-souza](https://discourse.julialang.org/u/de-souza)\
**Post date:** [October 4, 2022, 4:29pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/4 "2022-10-04T16:29:34Z")

</div>

Hi, unfortunately I haven’t found a simpler way so far.

In the line

```julia
names = :f1, :f2, :f3, :f4

```

have you replaced :f1, :f2… with the names of your variables? For example, if you set

```julia
@variables u(..) v(..) r(..) s(..)

```

the names should become

```julia
names = :u, :v, :r, :s

```

---

<div class="post-metadata">

**Author:** ![LeoCott](https://avatars.discourse-cdn.com/v4/letter/l/94ad74/32.png) [@LeoCott](https://discourse.julialang.org/u/LeoCott)\
**Post date:** [October 5, 2022, 8:14am UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/5 "2022-10-05T08:14:45Z")

</div>

Update @de-souza, i recompiled NeuralPDE@v5.3.0 and now it is working properly! Thanks a lot!

---

<div class="post-metadata">

**Author:** ![de-souza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/de-souza/32/43417_2.png) [@de-souza](https://discourse.julialang.org/u/de-souza)\
**Post date:** [October 5, 2022, 10:30am UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/6 "2022-10-05T10:30:51Z")

</div>

I’m glad it worked. I submitted a pull request with the fix to ComponentArrays.

> <https://github.com/jonniedie/ComponentArrays.jl/pull/166>
>
> Currently constructing ComponentArrays on the GPU with for example \`ComponentArr…ay(a=CUDA.ones(2,3), b=CUDA.ones(4,5))\` fails with an error.
> 
> \`\`\`
> ERROR: CuArray only supports element types that are stored inline
> \`\`\`
> 
> This fails because constructing the ComponentArray causes the creation of an array of arrays, which is not possible on the GPU. This happens in the \`pushcat!\` call at line 179 in componentarray.jl:
> 
> https://github.com/jonniedie/ComponentArrays.jl/blob/a7d6f7d126c5aeb711b65cf74eff8f8cbbeca5d0/src/componentarray.jl#L178-L182
> 
> This pull request fixes the issue by constructing the ComponentArray on the CPU and moving it back to the GPU afterwards.
> 
> Fixes #158. Needed for SciML/NeuralPDE.jl#594.