Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

GPU training does not work with at least Flux optimizers #141

Closed
AzamatB opened this issue Sep 5, 2020 · 12 comments · Fixed by #183
Closed

GPU training does not work with at least Flux optimizers #141

AzamatB opened this issue Sep 5, 2020 · 12 comments · Fixed by #183

Comments

@AzamatB
Copy link

AzamatB commented Sep 5, 2020

Consider an example from the docs demonstrating GPU training: https://neuralpde.sciml.ai/dev/examples/pinns_example/#Example-8-:-Fokker-Planck-equation-with-GPU-acceleration-1

If we call it with Flux's optimizers, for example, like GalacticOptim.solve(prob, ADAM()) then that dispatches to this implementation, which won't run on GPU. We can see it from

julia> _loss = function(θ)
           x = prob.f.f(θ, prob.p)
       end
SciML/GalacticOptim.jl#15 (generic function with 1 method)
julia> θ = copy(prob.x)   ### still on CPU ###
193-element Array{Float32,1}:
  ⋮
julia> gs = Flux.Zygote.gradient(ps = Flux.params(θ)) do
         x = _loss(θ)
         first(x)
       end
Grads(...)
julia> gs[θ]    ### still on CPU ###
193-element Array{Float64,1}:
@ChrisRackauckas
Copy link
Member

This doesn't seem like a GalacticOptim problem. NeuralPDE.jl should've created a problem where the parameters were on the GPU. Transferring.

@ChrisRackauckas ChrisRackauckas transferred this issue from SciML/Optimization.jl Sep 5, 2020
@AzamatB
Copy link
Author

AzamatB commented Sep 5, 2020 •
edited
Loading

So it looks like currently to train on GPU you have to explicitly initialize neural net parameters and move them to GPU and then pass those to PhysicsInformedNN like

chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) |> gpu
initθ = initial_params(chain) |> gpu
discretization = NeuralPDE.PhysicsInformedNN(dx, chain, initθ)

Here is the full reproducible example

using ModelingToolkit, NeuralPDE, Flux, DiffEqFlux, GalacticOptim, CUDA
CUDA.allowscalar(false)

# the example took from this article https://arxiv.org/abs/1910.10503
@parameters x θ
@variables p(..)
@derivatives Dx'~x
@derivatives Dxx''~x

#2D PDE
α = 0.3; β = 0.5; _σ = 0.5
# Discretization
dx = 0.05
# here we use normalization condition: dx*p(x,θ) ~ 1, in order to get non-zero solution.
eq  = [(α - 3*β*x^2)*p(x,θ) + (α*x - β*x^3)*Dx(p(x,θ)) ~ (_σ^2/2)*Dxx(p(x,θ)),
       dx*p(x,θ) ~ 1.]

# Initial and boundary conditions
bcs = [p(-2.2,θ) ~ 0. ,p(2.2,θ) ~ 0. , p(-2.2,θ) ~ p(2.2,θ)]

# Space and time domains
domains = [x ∈ IntervalDomain(-2.2,2.2)]

# Neural network
chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) |> gpu
initθ = initial_params(chain) |> gpu
discretization = NeuralPDE.PhysicsInformedNN(dx, chain, initθ)

pde_system = PDESystem(eq,bcs,domains,[x],[p])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM())

which now throws an error:

julia> res = GalacticOptim.solve(prob, ADAM())
ERROR: ArgumentError: cannot take the CPU address of a CuArray{Float32,1}
  Stacktrace:
   [1] unsafe_convert(::Type{Ptr{Float32}}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\CUDA\dZvbp\src\array.jl:237
   [2] unsafe_convert(::Type{Ptr{Float32}}, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}) at .\reshapedarray.jl:292
   [3] gemv!(::Char, ::Float32, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Float32, ::CuArray{Float32,1}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\blas.jl:626
Full trace julia> res = GalacticOptim.solve(prob, Flux.ADAM(); cb = cb) ERROR: ArgumentError: cannot take the CPU address of a CuArray{Float32,1} Stacktrace: [1] unsafe_convert(::Type{Ptr{Float32}}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\CUDA\dZvbp\src\array.jl:237 [2] unsafe_convert(::Type{Ptr{Float32}}, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}) at .\reshapedarray.jl:292 [3] gemv!(::Char, ::Float32, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Float32, ::CuArray{Float32,1}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\blas.jl:626 [4] gemv!(::CuArray{Float32,1}, ::Char, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Bool, ::Bool) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:470 [5] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:66 [inlined] [6] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:208 [inlined] [7] *(::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:47 [8] adjoint at D:\Users\Azamat\.julia\packages\DiffEqFlux\1xTu3\src\fast_layers.jl:57 [inlined] [9] _pullback at D:\Users\Azamat\.julia\packages\ZygoteRules\6nssF\src\adjoint.jl:47 [inlined] [10] applychain at D:\Users\Azamat\.julia\packages\DiffEqFlux\1xTu3\src\fast_layers.jl:20 [inlined] [11] _pullback(::Zygote.Context, ::typeof(DiffEqFlux.applychain), ::Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [12] FastChain at D:\Users\Azamat\.julia\packages\DiffEqFlux\1xTu3\src\fast_layers.jl:21 [inlined] [13] _pullback(::Zygote.Context, ::FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [14] #150 at D:\Users\Azamat\.julia\packages\NeuralPDE\Bp7SF\src\pinns_pde_solve.jl:330 [inlined] [15] _pullback(::Zygote.Context, ::NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}}, ::Array{Float64,1}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [16] #244 at D:\Users\Azamat\.julia\packages\NeuralPDE\Bp7SF\src\pinns_pde_solve.jl:247 [inlined] [17] _pullback(::Zygote.Context, ::NeuralPDE.var"#244#246"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}}}, ::Float64, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [18] #243 at D:\Users\Azamat\.julia\packages\NeuralPDE\Bp7SF\src\pinns_pde_solve.jl:262 [inlined] [19] _pullback(::Zygote.Context, ::NeuralPDE.var"#243#245", ::NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}}, ::Array{Float64,1}, ::CuArray{Float32,1}, ::NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [20] inner_loss at D:\Users\Azamat\.julia\packages\NeuralPDE\Bp7SF\src\pinns_pde_solve.jl:393 [inlined] [21] _pullback(::Zygote.Context, ::NeuralPDE.var"#inner_loss#169", ::NeuralPDE.var"#243#245", ::NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}}, ::Array{Float64,1}, ::CuArray{Float32,1}, ::NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [22] #168 at .\none:0 [inlined] [23] _pullback(::Zygote.Context, ::NeuralPDE.var"#168#173"{NeuralPDE.var"#243#245",CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"}, ::Array{Float64,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [24] MappingRF at .\reduce.jl:93 [inlined] [25] _pullback(::Zygote.Context, ::Base.MappingRF{NeuralPDE.var"#168#173"{NeuralPDE.var"#243#245",CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"},Base.MappingRF{typeof(abs2),Base.BottomRF{typeof(Base.add_sum)}}}, ::Base._InitialValue, ::Array{Float64,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [26] _foldl_impl at .\reduce.jl:58 [inlined] [27] _pullback(::Zygote.Context, ::typeof(Base._foldl_impl), ::Base.MappingRF{NeuralPDE.var"#168#173"{NeuralPDE.var"#243#245",CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"},Base.MappingRF{typeof(abs2),Base.BottomRF{typeof(Base.add_sum)}}}, ::Base._InitialValue, ::Array{Any,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [28] foldl_impl at .\reduce.jl:48 [inlined] [29] _pullback(::Zygote.Context, ::typeof(Base.foldl_impl), ::Base.MappingRF{NeuralPDE.var"#168#173"{NeuralPDE.var"#243#245",CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"},Base.MappingRF{typeof(abs2),Base.BottomRF{typeof(Base.add_sum)}}}, ::NamedTuple{(),Tuple{}}, ::Array{Any,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [30] mapfoldl_impl at .\reduce.jl:44 [inlined] [31] #mapfoldl#204 at .\reduce.jl:160 [inlined] [32] mapfoldl at .\reduce.jl:160 [inlined] [33] #mapreduce#208 at .\reduce.jl:287 [inlined] [34] mapreduce at .\reduce.jl:287 [inlined] [35] sum at .\reduce.jl:494 [inlined] [36] #167 at .\none:0 [inlined] [37] MappingRF at .\reduce.jl:93 [inlined] [38] _pullback(::Zygote.Context, ::Base.MappingRF{NeuralPDE.var"#167#172"{CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"},Base.BottomRF{typeof(Base.add_sum)}}, ::Base._InitialValue, ::Tuple{NeuralPDE.var"#243#245",Array{Any,1}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [39] _foldl_impl at .\reduce.jl:58 [inlined] [40] _pullback(::Zygote.Context, ::typeof(Base._foldl_impl), ::Base.MappingRF{NeuralPDE.var"#167#172"{CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"},Base.BottomRF{typeof(Base.add_sum)}}, ::Base._InitialValue, ::Base.Iterators.Zip{Tuple{Array{NeuralPDE.var"#243#245",1},Array{Array{Any,1},1}}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [41] foldl_impl at .\reduce.jl:48 [inlined] [42] _pullback(::Zygote.Context, ::typeof(Base.foldl_impl), ::Base.MappingRF{NeuralPDE.var"#167#172"{CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"},Base.BottomRF{typeof(Base.add_sum)}}, ::NamedTuple{(),Tuple{}}, ::Base.Iterators.Zip{Tuple{Array{NeuralPDE.var"#243#245",1},Array{Array{Any,1},1}}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [43] mapfoldl_impl at .\reduce.jl:44 [inlined] [44] #mapfoldl#204 at .\reduce.jl:160 [inlined] [45] mapfoldl at .\reduce.jl:160 [inlined] [46] #mapreduce#208 at .\reduce.jl:287 [inlined] [47] mapreduce at .\reduce.jl:287 [inlined] [48] sum at .\reduce.jl:494 [inlined] [49] sum at .\reduce.jl:511 [inlined] [50] _pullback(::Zygote.Context, ::typeof(sum), ::Base.Generator{Base.Iterators.Zip{Tuple{Array{NeuralPDE.var"#243#245",1},Array{Array{Any,1},1}}},NeuralPDE.var"#167#172"{CuArray{Float32,1},NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},NeuralPDE.var"#inner_loss#169"}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [51] #166 at D:\Users\Azamat\.julia\packages\NeuralPDE\Bp7SF\src\pinns_pde_solve.jl:428 [inlined] [52] _pullback(::Zygote.Context, ::NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [53] loss_function at D:\Users\Azamat\.julia\packages\NeuralPDE\Bp7SF\src\pinns_pde_solve.jl:485 [inlined] [54] _pullback(::Zygote.Context, ::NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}}, ::CuArray{Float32,1}, ::DiffEqBase.NullParameters) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [55] #9 at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:78 [inlined] [56] _pullback(::Zygote.Context, ::GalacticOptim.var"#9#16"{GalacticOptim.OptimizationProblem{OptimizationFunction{NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}},GalacticOptim.var"#94#100"{GalacticOptim.var"#93#99"{DiffEqBase.NullParameters,NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}}}},GalacticOptim.var"#95#101"{GalacticOptim.var"#93#99"{DiffEqBase.NullParameters,NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}}}},GalacticOptim.var"#98#104",Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}},CuArray{Float32,1},DiffEqBase.NullParameters,Nothing,Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}}}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [57] #12 at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:89 [inlined] [58] _pullback(::Zygote.Context, ::GalacticOptim.var"#12#19"{CuArray{Float32,1}}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0 [59] pullback(::Function, ::Zygote.Params) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface.jl:172 [60] gradient(::Function, ::Zygote.Params) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface.jl:53 [61] macro expansion at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:88 [inlined] [62] macro expansion at D:\Users\Azamat\.julia\packages\ProgressLogging\BBN0b\src\ProgressLogging.jl:328 [inlined] [63] (::GalacticOptim.var"#11#18"{var"#29#30",Int64,Bool,Bool,CuArray{Float32,1},Zygote.Params})() at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:54 [64] maybe_with_logger(::GalacticOptim.var"#11#18"{var"#29#30",Int64,Bool,Bool,CuArray{Float32,1},Zygote.Params}, ::Nothing) at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:29 [65] __solve(::GalacticOptim.OptimizationProblem{OptimizationFunction{NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}},GalacticOptim.var"#94#100"{GalacticOptim.var"#93#99"{DiffEqBase.NullParameters,NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}}}},GalacticOptim.var"#95#101"{GalacticOptim.var"#93#99"{DiffEqBase.NullParameters,NeuralPDE.var"#loss_function#178"{NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"},NeuralPDE.var"#166#171"{NeuralPDE.var"#150#154"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#159#163"{Array{Array{Float64,1},1}},Int64,NeuralPDE.var"#inner_loss#169"}}}},GalacticOptim.var"#98#104",Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}},CuArray{Float32,1},DiffEqBase.NullParameters,Nothing,Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}}, ::ADAM; cb::Function, maxiters::Int64, progress::Bool, save_best::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:53 [66] #solve#3 at D:\Users\Azamat\.julia\packages\GalacticOptim\8zvcf\src\solve.jl:3 [inlined]

@AzamatB
Copy link
Author

AzamatB commented Sep 5, 2020

I can prepare a PR fixing this if someone can explain what is the issue here and how to fix it

@ChrisRackauckas
Copy link
Member

You didn't call discretize in your example.

@AzamatB
Copy link
Author

AzamatB commented Sep 5, 2020 •
edited
Loading

Ah, I did call it, but it was a pasting typo. I've edited the example, thanks. The reported error still persists

@ChrisRackauckas
Copy link
Member

Are you on the latest version? It looks like it should just work.

@ChrisRackauckas
Copy link
Member

Oh, the changes are on master. Let me tag.

@AzamatB
Copy link
Author

AzamatB commented Sep 6, 2020

I'm on the master for both NeuralPDE and GalacticOptim:

(PoroelasticityPDESystem) pkg> status
Status `D:\Users\Azamat\Desktop\PoroelasticityPDESystem\Project.toml`
  [052768ef] CUDA v1.3.3
  [aae7a2af] DiffEqFlux v1.21.0
  [587475ba] Flux v0.11.1
  [a75be94c] GalacticOptim v0.1.3 `https://github.com/SciML/GalacticOptim.jl.git#master`
  [961ee093] ModelingToolkit v3.20.0
  [315f7962] NeuralPDE v2.2.0 `https://github.com/SciML/NeuralPDE.jl.git#master`
  [429524aa] Optim v0.22.0

julia> versioninfo()
Julia Version 1.5.1
Commit 697e782ab8 (2020-08-25 20:08 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i7-8700 CPU @ 3.20GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, skylake)
Environment:
  JULIA_EDITOR = "D:\Users\Azamat\AppData\Local\atom\app-1.50.0\atom.exe"  -a
  JULIA_NUM_THREADS = 6

Again, here is the example, which is the same as above (but without comments) and it still throws the error

using ModelingToolkit, NeuralPDE, Flux, DiffEqFlux, GalacticOptim, CUDA
CUDA.allowscalar(false)

@parameters x θ
@variables p(..)
@derivatives Dx'~x
@derivatives Dxx''~x

α = 0.3; β = 0.5; _σ = 0.5; dx = 0.05
eq  = [(α - 3*β*x^2)*p(x,θ) + (α*x - β*x^3)*Dx(p(x,θ)) ~ (_σ^2/2)*Dxx(p(x,θ)),
       dx*p(x,θ) ~ 1.]
bcs = [p(-2.2,θ) ~ 0. ,p(2.2,θ) ~ 0. , p(-2.2,θ) ~ p(2.2,θ)]

domains = [x ∈ IntervalDomain(-2.2,2.2)]

chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1)) |> gpu
initθ = initial_params(chain) |> gpu
discretization = NeuralPDE.PhysicsInformedNN(dx, chain, initθ)

pde_system = PDESystem(eq,bcs,domains,[x],[p])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM())

The last line throws

julia> res = GalacticOptim.solve(prob, ADAM())
ERROR: ArgumentError: cannot take the CPU address of a CuArray{Float32,1}
Stacktrace:
 [1] unsafe_convert(::Type{Ptr{Float32}}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\CUDA\dZvbp\src\array.jl:237
 [2] unsafe_convert(::Type{Ptr{Float32}}, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}) at .\reshapedarray.jl:292
 [3] gemv!(::Char, ::Float32, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Float32, ::CuArray{Float32,1}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\blas.jl:626
 [4] gemv!(::CuArray{Float32,1}, ::Char, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Bool, ::Bool) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:470
 [5] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:66 [inlined]
 [6] mul! at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:208 [inlined]
 [7] *(::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.5\LinearAlgebra\src\matmul.jl:47  
 [8] adjoint at D:\Users\Azamat\.julia\packages\DiffEqFlux\1xTu3\src\fast_layers.jl:57 [inlined]
 [9] _pullback at D:\Users\Azamat\.julia\packages\ZygoteRules\6nssF\src\adjoint.jl:47 [inlined]
 [10] applychain at D:\Users\Azamat\.julia\packages\DiffEqFlux\1xTu3\src\fast_layers.jl:20 [inlined]
 [11] _pullback(::Zygote.Context, ::typeof(DiffEqFlux.applychain), ::Tuple{FastDense{typeof(σ), DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0
 [12] FastChain at D:\Users\Azamat\.julia\packages\DiffEqFlux\1xTu3\src\fast_layers.jl:21 [inlined]
 [13] _pullback(::Zygote.Context, ::FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#115"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at D:\Users\Azamat\.julia\packages\Zygote\Xgcgs\src\compiler\interface2.jl:0
 ...

@ChrisRackauckas
Copy link
Member

Yes, the current GPU tests are not actually on the GPU (@KirillZubov). The issue is the training data. The training dataset generates a bunch of CPU-bond arrays. There's a few issues with that. For one, there's no reason to generate the arrays: with that training strategy they can be left as lazy arrays. That's necessary because those are the biggest use of memory, so the current strategy won't even work on the GPU. The product iterator should just never create an array.

@KirillZubov
Copy link
Member

@AzamatB thanks for point out this bug.
@ChrisRackauckas ok, I try to fix it this week.

@AzamatB
Copy link
Author

AzamatB commented Sep 14, 2020

Bump. Would be great to resolve this.

@AzamatB
Copy link
Author

AzamatB commented Nov 16, 2020

Bump.

@KirillZubov KirillZubov linked a pull request Dec 2, 2020 that will close this issue
[WIP] GPU pinns #183
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

[WIP] GPU pinns SciML/NeuralPDE.jl
3 participants
@ChrisRackauckas @AzamatB @KirillZubov