the example of Bayesian Physics informed Neural Network ODEs Solvers is nice.
after executing, it takes about 9 minutes for one inference.
I tried to use cuda gpu with the help of Monica gpt, but failed to execute.
can anyone give suggestions ? or share another good examples.
thanks.
using NeuralPDE, Lux, Plots, OrdinaryDiffEq, Distributions, Random
using LuxCUDA,CUDA
if !CUDA.has_cuda()
error(“CUDA is not available. Please ensure that you have a compatible GPU and the CUDA toolkit installed.”)
end
function lotka_volterra(u, p, t)
# Model parameters.
α, β, γ, δ = p
# Current state.
x, y = u
# Evaluate differential equations.
dx = (α - β * y) * x # prey
dy = (δ * x - γ) * y # predator
return CUDA.CuArray([dx, dy])
end
initial-value problem.
u0 = CUDA.CuArray(Float32[1.0, 1.0]) # Create CuArray directly
p = CUDA.CuArray(Float32[1.5, 1.0, 3.0, 1.0]) # Create CuArray directly
tspan = (0.0f0, 4.0f0)
prob = ODEProblem(lotka_volterra, u0, tspan, p)
Solve using OrdinaryDiffEq.jl solver
dt = 0.01f0
solution = solve(prob, Tsit5(); saveat = dt,reltol=1e-6, abstol=1e-6)
Dataset creation for parameter estimation (30% noise)
time = Float32.(solution.t)
u = hcat(Float32.(solution.u)…)
x = u[1, :] .+ (u[1, :]) .* (0.3 .* randn(length(u[1, :])))
y = u[2, :] .+ (u[2, :]) .* (0.3 .* randn(length(u[2, :])))
Move datasets to GPU
x_gpu = CUDA.CuArray(Float32.(x)) # Create CuArray directly
y_gpu = CUDA.CuArray(Float32.(y)) # Create CuArray directly
time_gpu = CUDA.CuArray(Float32.(time)); # Create CuArray directly
dataset = [x_gpu, y_gpu, time_gpu]; # Create dataset on GPU
Plotting the noisy data
plot(time, x, label = “Noisy Prey (x)”, xlabel=“Time”, ylabel=“Population”, title=“Lotka-Volterra Model”)
plot!(time, y, label = “Noisy Predator (y)”)
plot!(solution.t, solution.u[1, :], label = “True Prey (x)”, linestyle=:dash)
plot!(solution.t, solution.u[2, :], label = “True Predator (y)”, linestyle=:dash)
Neural Networks must have 2 outputs as u → [dx,dy] in function lotka_volterra()
chain = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh),Lux.Dense(6, 2))|> CUDA.Chain # Move the chain to GPU
Setup for Bayesian Neural Network for parameter estimation
alg = BNNODE(chain;
dataset = dataset,
draw_samples = 1000,
l2std = [0.1f0, 0.1f0],
phystd = [0.1f0, 0.1f0],
priorsNNw = (0.0f0, 3.0f0),
param = [
Normal(1f0, 2f0),
Normal(2f0, 2f0),
Normal(2f0, 2f0),
Normal(0f0, 2f0)
],
progress = false
)
sol_pestim = solve(prob, alg; saveat = dt)
Plotting the estimated results
estimated_u = hcat(sol_pestim.u…) # Combine estimated solution arrays
estimated_x = estimated_u[1, :]
estimated_y = estimated_u[2, :]
Final plot of the results
plot(time, estimated_x, label = “Estimated Prey (x)”, xlabel=“Time”, ylabel=“Population”, title=“Lotka-Volterra Model - Parameter Estimation”)
plot!(time, estimated_y, label = “Estimated Predator (y)”)
plot!(solution.t, solution.u[1, :], label = “True Prey (x)”, linestyle=:dash)
plot!(solution.t, solution.u[2, :], label = “True Predator (y)”, linestyle=:dash)
plot!(time, x, label = “Noisy Prey (x)”, linestyle=:dot)
plot!(time, y, label = “Noisy Predator (y)”, linestyle=:dot)
===
the error is:
Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations do not execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.