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

Specific Domains GPU
1

Hi all,

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

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

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

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:

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!

2

Hi, I faced a similar problem. Somebody else also created an issue on GitHub.

Here is how I fixed it:

  1. Define all the initial parameters as a single ComponentArray.
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 (necessary for me as of NeuralPDE@v5.3.0 and ComponentArrays@v0.13.2).
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.
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.

2 Likes
3

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 ?

4

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

In the line

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

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

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

the names should become

names = :u, :v, :r, :s
5

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

6

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