PINN using Flux

[andres]

Hello, folks

I’m trying to create a Physics-Informed Neural Network (PINN) to solve the harmonic oscillator using Flux.

The idea was to make a Julia version based on the following link, whose implementation is in Python:

GitHub - benmoseley/harmonic-oscillator-pinn: Code accompanying my blog post:...

The first part of the program is to create a traditional neural network, and that part is okay.

begin
    using Flux
    using LinearAlgebra
    using Plots
    using Statistics
    using Zygote
    gr()
end

function ocilator(δ, ω0, x)
    @assert δ < ω0
    ω = √(ω0^2 - δ^2)
    ϕ = atan(-δ / ω)
    A = 1 / (2 * cos(ϕ))
    y = 2 * A * exp.(-δ .* x) .* cos.(ω * x .+ ϕ)
end

function nn(N_INPUT::Int=1, N_OUTPUT::Int=1, N_HIDDEN::Int=32, N_LAYERS::Int=3)
    head = Dense(N_INPUT => N_HIDDEN, tanh)
    hidden = Chain(
        [Dense(N_HIDDEN => N_HIDDEN, tanh) for _ in 1:N_LAYERS]
    )
    tail = Dense(N_HIDDEN => N_OUTPUT)

    return Chain(head, hidden, tail)
end

function plot_result(epoch, x, y, x_data, y_data, yh, xp=nothing)
    p = plot(size=(16, 9) .* 40, axis=:off, xlims=(-0.05, 1.05), ylims=(-1.1, 1.1))
    plot!(p, x[:], y[:], color="grey", lw=2, alpha=0.7, label="Exact solution")
    plot!(p, x[:], yh[:], color="blue", lw=4, alpha=0.8, label="Neural network prediction", xlims=(-0.05, 1.05), ylims=(-1.1, 1.1))
    scatter!(p, x_data[:], y_data[:], color="orange", alpha=0.4, label="Training data",)
    if !isnothing(xp)
        scatter!(p, xp[:], zeros(size(xp))[:], color="green", alpha=0.4, label="Physics loss training locations",)
    end

    xlims!(p, (-0.05, 1.05))
    ylims!(p, (-1.1, 1.1))
    plot(p, title="Época $epoch")
end

begin
    δ, ω0 = 2, 20
    x = Float32.(hcat(LinRange(0, 1, 500)...)) |> gpu
    y = ocilator(δ, ω0, x)

    x_data = x[:, 1:20:200] |> gpu
    y_data = y[:, 1:20:200] |> gpu

    yh = (nn() |> gpu)(x)

    p = plot(x[:], y[:], label="Solução exata")
    scatter!(p, x_data[:], y_data[:], label="Dados de treino")
    plot(p, size=(16, 9) .* 40, axis=:off)
end

begin
    plot_result(0, x, y, x_data, y_data, yh)
end

begin
    model = nn() |> gpu
    opt = Flux.setup(Adam(), model)
    epochs = 1000
    anim = Animation()

    for epoch in 1:epochs
        grads = Flux.gradient(model) do m
            y_hat = m(x_data)
            Flux.mse(y_hat, y_data)
        end
        Flux.update!(opt, model, grads[1])

        if epoch % 10 == 0
            y_hat = model(x)
            frame(anim, plot_result(epoch, x, y, x_data, y_data, y_hat))
        end
    end
    gif(anim)
end

The second part is to use the differential equation in the objective function and thus achieve a better approximation.

However, I am facing all possible issues, and it keeps throwing errors.

begin
    xp = hcat(Float32.(0:1/32:1)...) |> gpu
    μ, κ = 2 * δ, ω0^2

    model = nn() |> gpu
    opt = Flux.setup(Adam(), model)
    epochs = 10
    anim = Animation()

    for epoch in 1:epochs
        val, grads = Flux.withgradient(model) do m
            y_hat = m(x_data)
            ls1 = Flux.mse(y_hat, y_data)

            yhp = m(xp)
            dx = Zygote.gradient(x -> sum(m(x)), xp)[1]
            dx2 = Zygote.diaghessian(x -> sum(m(x)), xp)[1]
            p = dx2 + μ * dx + κ * yhp
            ls2 = (1e-4) * mean(p .^ 2)

            ls1 + ls2
        end

        Flux.update!(opt, model, grads[1])
    end
end

I am calculating the gradient and the Hessian to obtain the solution to the differential equation, and it is precisely in this part that I am encountering the errors.

I’ve tried calculating it on the GPU, on the CPU, everywhere, but I have no idea how to proceed.

[tomaklutfu]

From the error, it seems diaghessian uses setindex! somewhere; hence, the error in finding gradient with respect to model parameters. It seems hard to overcome some issues when using 2nd derivatives as described in ref (Limitations · Zygote).

[tomaklutfu]

Also, the reference suggests using another package like ForwardDiff over Zygote.gradient so this seems to work(at least no error)

val, grads = Flux.withgradient(model) do m
  y_hat = m(x_data)
  ls1 = Flux.mse(y_hat, y_data)
  
  yhp = m(xp)
  dx  = map(y->ForwardDiff.derivative(x -> sum(m([x])), y), xp)
  dx2 = map(z->ForwardDiff.derivative(y->ForwardDiff.derivative(x -> sum(m([x])), y), z), xp)
  p = κ * yhp + μ * dx + dx2
  ls2 = (1e-4) * mean(p.^2)
  
  ls1 + ls2
  end

However, I cannot make it learn better because derivatives become too noisy, I think. It was my mistake in code model should be m.

[andres]

Thank you for your response.

When using the CPU, I am receiving the following warning, and the solution is not so great:

┌ Warning: `ForwardDiff.derivative(f, x)` within Zygote cannot track gradients with respect to `f`,
│ and `f` appears to be a closure, or a struct with fields (according to `issingletontype(typeof(f))`).
│ typeof(f) = var"#195#201"{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}, Chain{Vector{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}
└ @ Zygote ~/.julia/packages/Zygote/YYT6v/src/lib/forward.jl:158

┌ Warning: `ForwardDiff.derivative(f, x)` within Zygote cannot track gradients with respect to `f`,
│ and `f` appears to be a closure, or a struct with fields (according to `issingletontype(typeof(f))`).
│ typeof(f) = var"#197#203"{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}, Chain{Vector{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}
└ @ Zygote ~/.julia/packages/Zygote/YYT6v/src/lib/forward.jl:158

When using the GPU, I am getting the following error, and nothing is working:

ERROR: GPU compilation of MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::CUDA.CuKernelContext, ::CuDeviceMatrix{Tuple{Float32, Zygote.Pullback{Tuple{var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}, Float32}, Tuple{Zygote.var"#2160#back#297"{Zygote.var"#back#296"{:m, Zygote.Context{false}, var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}, Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}}, ComposedFunction{Zygote.var"#4206#back#1471"{Zygote.var"#1467#1470"{Float32, Matrix{Float32}}}, typeof(ZygoteRules.unthunk_tangent)}, Zygote.var"#2187#back#307"{Zygote.Jnew{var"#35#41"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}, Nothing, false}}, Zygote.var"#2160#back#297"{Zygote.var"#back#296"{:m, Zygote.Context{false}, var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}, Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}}}}}, 1}, ::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{2}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, Zygote.var"#655#659"{Zygote.Context{false}, var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}}, Tuple{Base.Broadcast.Extruded{CuDeviceMatrix{Float32, 1}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}}}, ::Int64) failed
KernelError: passing and using non-bitstype argument

Argument 4 to your kernel function is of type Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{2}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, Zygote.var"#655#659"{Zygote.Context{false}, var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}}, Tuple{Base.Broadcast.Extruded{CuDeviceMatrix{Float32, 1}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}}}, which is not isbits:
  .f is of type Zygote.var"#655#659"{Zygote.Context{false}, var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}}} which is not isbits.
    .cx is of type Zygote.Context{false} which is not isbits.
      .cache is of type Union{Nothing, IdDict{Any, Any}} which is not isbits.
    .f is of type var"#34#40"{Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}} which is not isbits.
      .m is of type Chain{Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}} which is not isbits.
        .layers is of type Tuple{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}} which is not isbits.
          .1 is of type Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}} which is not isbits.
            .weight is of type CuArray{Float32, 2, CUDA.Mem.DeviceBuffer} which is not isbits.
              .storage is of type Union{Nothing, CUDA.ArrayStorage{CUDA.Mem.DeviceBuffer}} which is not isbits.
            .bias is of type CuArray{Float32, 1, CUDA.Mem.DeviceBuffer} which is not isbits.
              .storage is of type Union{Nothing, CUDA.ArrayStorage{CUDA.Mem.DeviceBuffer}} which is not isbits.
          .2 is of type Chain{Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}} which is not isbits.
            .layers is of type Vector{Dense{typeof(tanh), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}} which is not isbits.
          .3 is of type Dense{typeof(identity), CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}} which is not isbits.
            .weight is of type CuArray{Float32, 2, CUDA.Mem.DeviceBuffer} which is not isbits.
              .storage is of type Union{Nothing, CUDA.ArrayStorage{CUDA.Mem.DeviceBuffer}} which is not isbits.
            .bias is of type CuArray{Float32, 1, CUDA.Mem.DeviceBuffer} which is not isbits.
              .storage is of type Union{Nothing, CUDA.ArrayStorage{CUDA.Mem.DeviceBuffer}} which is not isbits.

I will try to search for other libraries to perform this anomaly detection in a way that doesn’t result in as many errors.

Thanks again.

[ChrisRackauckas]

In this day and age I’d try using Enzyme.jl on Lux.jl for this. Flux might give you issues. NeuralPDE.jl is a library for PINNs and we just finished a complete port over to Lux.jl in order to improve robustness and performance.