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

2 Likes

What ljne is erroring exactly (on mobile and can’t run it
)? Before subsetting you will need to pass it back to the CPU to not get the scalar error, so I think the issue is when you get the data out of `estimated_u`.

This won’t work for GPU Arrays. (x, y) is doing a scalar getindex to fetch data from `u`. You need to rewrite that with broadcasting and such. See Neural ODEs on GPUs · DiffEqFlux.jl for a different example on how to get these to work on GPUs

2 Likes