# Saving/loading a neural network trained with NeuralPDE

**URL:** <https://discourse.julialang.org/t/saving-loading-a-neural-network-trained-with-neuralpde/117583>\
**Category:** General Usage\
**Tags:** jld2, neural-network, lux\
**Created:** [July 29, 2024, 8:48am UTC](https://discourse.julialang.org/t/saving-loading-a-neural-network-trained-with-neuralpde/117583 "2024-07-29T08:48:22Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![gradureau](https://avatars.discourse-cdn.com/v4/letter/g/67e7ee/32.png) [@gradureau](https://discourse.julialang.org/u/gradureau)\
**Post date:** [July 29, 2024, 8:48am UTC](https://discourse.julialang.org/t/saving-loading-a-neural-network-trained-with-neuralpde/117583/1 "2024-07-29T08:48:22Z")

</div>

Hello everybody,

I am using neuralPDE to train neural networks and I wish to save them in files to use them later. It works just fine if I use the standard Lux layers. Yet I wish to use the following custom layer at the end of the model, in order to inforce the initial conditions in the architecture of the neural network:

```julia
   struct ICLayer_1D <: Lux.AbstractExplicitContainerLayer{(:model,)}
        model # Model of the layer
        IC_fct # Initial condition function
        t0 # Initial time
        layers # Layers of the model
    end

    function ICLayer_1D(model, IC_fct, t0)
        # Constructor of the layer
        return ICLayer_1D(model, IC_fct, t0, model.layers)
    end

    function (n::ICLayer_1D)(input, ps, st)
        # Operation of the layer
        # The model model must have 2 inputs and 3 outputs.
        # IC_fct must have 1 input.

        # Get model's inputs
        t = input[1,:]
        x = input[2,:]

        # Get model's outputs
        output, st_new = n.model(input, ps, st)

        # Expression of the layers inputs
        α = (t .- n.t0) .* output[1,:] .+ 1.0
        ξ_x = (t .- n.t0) .* output[2,:] .+ x
        β = (t .- n.t0) .* output[3,:]

        # Output of the layer
        return α.*n.IC_fct.(ξ_x) .+ β, st_new
    end

```

I have saved the parameters and the model is a JLD2 file and when I try to reconstruct the model using the command

`ps, st = Lux.setup(Random.default_rng(), nn)`

where nn is the loaded neural network, I get the following error:

```julia
ERROR: LoadError: MethodError: no method matching initialparameters(::TaskLocalRNG, ::JLD2.ReconstructedMutable{:ICLayer_1D, (:model, :IC_fct, :t0, :layers), NTuple{4, Any}})

Closest candidates are:
  initialparameters(::AbstractRNG, ::Any)
   @ LuxCore ~/.julia/packages/LuxCore/biwfu/src/LuxCore.jl:66
  initialparameters(::AbstractRNG, ::InstanceNorm)
   @ Lux ~/.julia/packages/Lux/JXc6P/src/layers/normalize.jl:354
  initialparameters(::AbstractRNG, ::NeuralPDE.dgm_lstm_layer)
   @ NeuralPDE ~/.julia/packages/NeuralPDE/IUP3N/src/dgm.jl:18
  ...

```

Would anyone know how to fix this?

Thank you for your answers!

---

<div class="post-metadata">

**Author:** ![JonasIsensee](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jonasisensee/32/4704_2.png) [@JonasIsensee](https://discourse.julialang.org/u/JonasIsensee)\
**Post date:** [July 29, 2024, 9:59am UTC](https://discourse.julialang.org/t/saving-loading-a-neural-network-trained-with-neuralpde/117583/2 "2024-07-29T09:59:28Z")

</div>

Please try to define your custom layer struct in the new session before attempting to load the file.

---

<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:** [July 30, 2024, 2:28am UTC](https://discourse.julialang.org/t/saving-loading-a-neural-network-trained-with-neuralpde/117583/3 "2024-07-30T02:28:22Z")

</div>

In general, it is not recommended to serialize Lux models. We guarantee that saved parameters and states can be loaded if there isn’t a breaking change. The same guarantee doesn’t hold for saving the architecture.

But in this specific scenario, `@JonasIsensee`’s suggestion will probably work

---

<div class="post-metadata">

**Author:** ![gradureau](https://avatars.discourse-cdn.com/v4/letter/g/67e7ee/32.png) [@gradureau](https://discourse.julialang.org/u/gradureau)\
**Post date:** [July 30, 2024, 4:42pm UTC](https://discourse.julialang.org/t/saving-loading-a-neural-network-trained-with-neuralpde/117583/4 "2024-07-30T16:42:39Z")

</div>

Perfect, It indeed works.