PINN using Flux

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:

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)

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
epochs = 1000
anim = Animation()

for epoch in 1:epochs
y_hat = m(x_data)
Flux.mse(y_hat, y_data)
end

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
epochs = 10
anim = Animation()

for epoch in 1:epochs
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

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.

1 Like

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).

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`.

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.

1 Like

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.