implement HybridPointSolver on gpu

Author: bgctw

dev/doubleMM.jl

@@ -11,49 +11,51 @@ using UnicodePlots
 using SimpleChains
 using Flux
 using MLUtils
+using CUDA
 
 rng = StableRNG(114)
 scenario = NTuple{0, Symbol}()
-#scenario = (:use_Flux,)
+scenario = (:use_Flux,)
 
 #------ setup synthetic data and training data loader
 (; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc
 ) = gen_hybridcase_synthetic(rng, DoubleMM.DoubleMMCase(); scenario);
+xM_cpu = xM
+if :use_Flux ∈ scenario
+    xM = CuArray(xM_cpu)
+end
 get_train_loader = (rng; n_batch, kwargs...) -> MLUtils.DataLoader((xM, xP, y_o, y_unc), batchsize = n_batch)
 σ_o = exp(first(y_unc)/2)
 
 # assign the train_loader, otherwise it eatch time creates another version of synthetic data
-prob0 = update(HybridProblem(DoubleMM.DoubleMMCase(); scenario); get_train_loader)
+prob0 = HVI.update(HybridProblem(DoubleMM.DoubleMMCase(); scenario); get_train_loader)
 
 #------- pointwise hybrid model fit
 #solver = HybridPointSolver(; alg = Adam(0.02), n_batch = 30)
 solver = HybridPointSolver(; alg = Adam(0.01), n_batch = 10)
 #solver = HybridPointSolver(; alg = Adam(), n_batch = 200)
 (; ϕ, resopt) = solve(prob0, solver; scenario,
     rng, callback = callback_loss(100), maxiters = 1200);
-prob0o = update(prob0; ϕg=ϕ.ϕg, θP=ϕ.θP)
-y_pred_global, y_pred, θMs = gf(prob0o, xM, xP);
+# update the problem with optimized parameters
+prob0o = HVI.update(prob0; ϕg=cpu_ca(ϕ).ϕg, θP=cpu_ca(ϕ).θP)
+y_pred_global, y_pred, θMs = gf(prob0o, xM, xP; scenario);
 scatterplot(θMs_true[1,:], θMs[1,:])
 scatterplot(θMs_true[2,:], θMs[2,:])
 
 # do a few steps without minibatching, 
 #   by providing the data rather than the DataLoader
-# train_loader0 = get_hybridproblem_train_dataloader(rng, prob0; scenario, n_batch=1000)
-# get_train_loader_data = (args...; kwargs...) -> train_loader0.data
-# prob1 = update(prob0o; get_train_loader = get_train_loader_data)
-prob1 = prob0o
-
-#solver1 = HybridPointSolver(; alg = Adam(0.05), n_batch = n_site)
 solver1 = HybridPointSolver(; alg = Adam(0.01), n_batch = n_site)
-(; ϕ, resopt) = solve(prob1, solver1; scenario, rng, 
+(; ϕ, resopt) = solve(prob0o, solver1; scenario, rng, 
     callback = callback_loss(20), maxiters = 600);
-prob1o = update(prob1; ϕg=ϕ.ϕg, θP=ϕ.θP)
-y_pred_global, y_pred, θMs = gf(prob1o, xM, xP);
+prob1o = HVI.update(prob0o; ϕg=cpu_ca(ϕ).ϕg, θP=cpu_ca(ϕ).θP);
+y_pred_global, y_pred, θMs = gf(prob1o, xM, xP; scenario);
 scatterplot(θMs_true[1,:], θMs[1,:])
 scatterplot(θMs_true[2,:], θMs[2,:])
 prob1o.θP
 scatterplot(vec(y_true), vec(y_pred))
 
+# still overestimating θMs
+
 () -> begin # with more iterations?
     prob2 = prob1o
     (; ϕ, resopt) = solve(prob2, solver1; scenario, rng, 
@@ -63,50 +65,55 @@ scatterplot(vec(y_true), vec(y_pred))
     prob2o.θP
 end
 
-#----------- fit g to true θMs 
-# and fit gf starting from true parameters
-prob = prob0
-g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario);
-(; transP, transM) = get_hybridproblem_transforms(prob; scenario)
-
-function loss_g(ϕg, x, g, transM)
-    ζMs = g(x, ϕg) # predict the log of the parameters
-    θMs = reduce(hcat, map(transM, eachcol(ζMs))) # transform each column
-    loss = sum(abs2, θMs .- θMs_true)
-    return loss, θMs
-end
-loss_g(ϕg0, xM, g, transM)
 
-optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g, transM)[1],
-    Optimization.AutoZygote())
-optprob = Optimization.OptimizationProblem(optf, ϕg0);
-res = Optimization.solve(optprob, Adam(0.015), callback = callback_loss(100), maxiters = 2000);
-
-ϕg_opt1 = res.u;
-l1, θMs = loss_g(ϕg_opt1, xM, g, transM)
-#scatterplot(θMs_true[1,:], θMs[1,:])
-scatterplot(θMs_true[2,:], θMs[2,:]) # able to fit θMs[2,:]
-
-prob3 = update(prob0, ϕg = ϕg_opt1, θP = θP_true)
-solver1 = HybridPointSolver(; alg = Adam(0.01), n_batch = n_site)
-(; ϕ, resopt) = solve(prob3, solver1; scenario, rng, 
-    callback = callback_loss(50), maxiters = 600);
-prob3o = update(prob3; ϕg=ϕ.ϕg, θP=ϕ.θP)
-y_pred_global, y_pred, θMs = gf(prob3o, xM, xP);
-scatterplot(θMs_true[2,:], θMs[2,:])
-prob3o.θP
-scatterplot(vec(y_true), vec(y_pred))
-scatterplot(vec(y_true), vec(y_o))
-scatterplot(vec(y_pred), vec(y_o))
+#----------- fit g to true θMs 
+() -> begin
+    # and fit gf starting from true parameters
+    prob = prob0
+    g, ϕg0_cpu = get_hybridproblem_MLapplicator(prob; scenario);
+    ϕg0 = (:use_Flux ∈ scenario) ? CuArray(ϕg0_cpu) : ϕg0_cpu
+    (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
+
+    function loss_g(ϕg, x, g, transM; gpu_handler = HVI.default_GPU_DataHandler)
+        ζMs = g(x, ϕg) # predict the log of the parameters
+        ζMs_cpu = gpu_handler(ζMs)
+        θMs = reduce(hcat, map(transM, eachcol(ζMs_cpu))) # transform each column
+        loss = sum(abs2, θMs .- θMs_true)
+        return loss, θMs
+    end
+    loss_g(ϕg0, xM, g, transM)
+
+    optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g, transM)[1],
+        Optimization.AutoZygote())
+    optprob = Optimization.OptimizationProblem(optf, ϕg0);
+    res = Optimization.solve(optprob, Adam(0.015), callback = callback_loss(100), maxiters = 2000);
+
+    ϕg_opt1 = res.u;
+    l1, θMs = loss_g(ϕg_opt1, xM, g, transM)
+    #scatterplot(θMs_true[1,:], θMs[1,:])
+    scatterplot(θMs_true[2,:], θMs[2,:]) # able to fit θMs[2,:]
+
+    prob3 = HVI.update(prob0, ϕg = Array(ϕg_opt1), θP = θP_true)
+    solver1 = HybridPointSolver(; alg = Adam(0.01), n_batch = n_site)
+    (; ϕ, resopt) = solve(prob3, solver1; scenario, rng, 
+        callback = callback_loss(50), maxiters = 600);
+    prob3o = HVI.update(prob3; ϕg=cpu_ca(ϕ).ϕg, θP=cpu_ca(ϕ).θP)
+    y_pred_global, y_pred, θMs = gf(prob3o, xM, xP; scenario);
+    scatterplot(θMs_true[2,:], θMs[2,:])
+    prob3o.θP
+    scatterplot(vec(y_true), vec(y_pred))
+    scatterplot(vec(y_true), vec(y_o))
+    scatterplot(vec(y_pred), vec(y_o))
 
-() -> begin # optimized loss is indeed lower than with true parameters
-    int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
-        ϕg = 1:length(prob0.ϕg), θP = prob0.θP))
-    loss_gf = get_loss_gf(prob0.g, prob0.transM, prob0.f, Float32[], int_ϕθP)
-    loss_gf(vcat(prob3.ϕg, prob3.θP), xM, xP, y_o, y_unc)[1]
-    loss_gf(vcat(prob3o.ϕg, prob3o.θP), xM, xP, y_o, y_unc)[1]
-    #
-    loss_gf(vcat(prob2o.ϕg, prob2o.θP), xM, xP, y_o, y_unc)[1]
+    () -> begin # optimized loss is indeed lower than with true parameters
+        int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
+            ϕg = 1:length(prob0.ϕg), θP = prob0.θP))
+        loss_gf = get_loss_gf(prob0.g, prob0.transM, prob0.f, Float32[], int_ϕθP)
+        loss_gf(vcat(prob3.ϕg, prob3.θP), xM, xP, y_o, y_unc)[1]
+        loss_gf(vcat(prob3o.ϕg, prob3o.θP), xM, xP, y_o, y_unc)[1]
+        #
+        loss_gf(vcat(prob2o.ϕg, prob2o.θP), xM, xP, y_o, y_unc)[1]
+    end
 end
 
 #----------- Hybrid Variational inference 

ext/HybridVariationalInferenceFluxExt.jl

@@ -55,6 +55,11 @@ function HVI.construct_3layer_MLApplicator(
     construct_ChainsApplicator(rng, g_chain, float_type)
 end
 
+function HVI.cpu_ca(ca::CA.ComponentArray)
+    CA.ComponentArray(cpu(CA.getdata(ca)), CA.getaxes(ca))
+end
+
+
 
 
 end # module

src/DoubleMM/f_doubleMM.jl

@@ -22,6 +22,12 @@ function HVI.get_hybridproblem_par_templates(::DoubleMMCase; scenario::NTuple =
     (; θP, θM)
 end
 
+function HVI.get_hybridproblem_MLapplicator(
+    rng::AbstractRNG, prob::HVI.DoubleMM.DoubleMMCase; scenario = ())
+    ml_engine = select_ml_engine(; scenario)
+    construct_3layer_MLApplicator(rng, prob, ml_engine; scenario)
+end
+
 function HVI.get_hybridproblem_transforms(::DoubleMMCase; scenario::NTuple = ())
     (; transP, transM)
 end
@@ -91,11 +97,6 @@ function HVI.gen_hybridcase_synthetic(rng::AbstractRNG, prob::DoubleMMCase;
     )
 end
 
-function HVI.get_hybridproblem_MLapplicator(
-    rng::AbstractRNG, prob::HVI.DoubleMM.DoubleMMCase; scenario = ())
-    ml_engine = select_ml_engine(; scenario)
-    construct_3layer_MLApplicator(rng, prob, ml_engine; scenario)
-end
 
 
 

src/HybridSolver.jl

@@ -24,7 +24,8 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
     f = get_hybridproblem_PBmodel(prob; scenario)
     y_global_o = FT[] # TODO
     loss_gf = get_loss_gf(g, transM, f, y_global_o, int_ϕθP)
-    #l1 = loss_gf(p0, train_loader...)[1]
+    # data1 = first(train_loader)
+    # l1 = loss_gf(p0, first(train_loader)...)[1]
     # Zygote.gradient(p0 -> loss_gf(p0, data1...)[1], p0)
     optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
         Optimization.AutoZygote())

src/HybridVariationalInference.jl

@@ -31,7 +31,7 @@ export AbstractHybridProblem, get_hybridproblem_MLapplicator, get_hybridproblem_
        get_hybridproblem_par_templates, get_hybridproblem_transforms, get_hybridproblem_train_dataloader,
        get_hybridproblem_neg_logden_obs, 
        get_hybridproblem_n_covar, 
-       update,
+       #update,
        gen_cov_pred
 include("AbstractHybridProblem.jl")
 
@@ -47,6 +47,9 @@ include("gencovar.jl")
 export callback_loss
 include("util_opt.jl")
 
+export cpu_ca
+include("util_ca.jl")
+
 export neg_logden_indep_normal, entropy_MvNormal
 include("logden_normal.jl")
 

src/gf.jl

@@ -24,8 +24,13 @@ function gf(g, transM, f, xM, xP, ϕg, θP; gpu_handler = default_GPU_DataHandle
     # @show first(ϕg,5)
     ζMs = g(xM, ϕg) # predict the log of the parameters
     ζMs_cpu = gpu_handler(ζMs)
+    if θP isa SubArray && !(gpu_handler isa NullGPUDataHandler) 
+        # otherwise Zyote fails on gpu_handler
+        θP = copy(θP)
+    end
+    θP_cpu = gpu_handler(CA.getdata(θP))
     θMs = reduce(hcat, map(transM, eachcol(ζMs_cpu))) # transform each column
-    y_pred_global, y_pred = f(θP, θMs, xP)
+    y_pred_global, y_pred = f(θP_cpu, θMs, xP)
     return y_pred_global, y_pred, θMs
 end
 
@@ -34,7 +39,8 @@ function gf(prob::AbstractHybridProblem, xM, xP, args...; scenario = (), kwargs.
     f = get_hybridproblem_PBmodel(prob; scenario)
     (; θP, θM) = get_hybridproblem_par_templates(prob; scenario)
     (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
-    gf(g, transM, f, xM, xP, ϕg, θP; kwargs...)
+    ϕg_dev, θP_dev = (:use_Flux ∈ scenario) ? (CuArray(ϕg), CuArray(CA.getdata(θP))) : (ϕg, CA.getdata(θP))
+    gf(g, transM, f, xM, xP, ϕg_dev, θP_dev; kwargs...)
 end
 
 """
@@ -50,7 +56,8 @@ function get_loss_gf(g, transM, f, y_o_global, int_ϕθP::AbstractComponentArray
         function loss_gf(p, xM, xP, y_o, y_unc)
             σ = exp.(y_unc ./ 2)
             pc = int_ϕθP(p)
-            y_pred_global, y_pred, θMs = gf(g, transM, f, xM, xP, pc.ϕg, pc.θP)
+            y_pred_global, y_pred, θMs = gf(
+                g, transM, f, xM, xP, CA.getdata(pc.ϕg), CA.getdata(pc.θP))
             loss = sum(abs2, (y_pred .- y_o) ./ σ) + sum(abs2, y_pred_global .- y_o_global)
             return loss, y_pred_global, y_pred, θMs
         end

src/util_ca.jl

@@ -0,0 +1,9 @@
+"""
+    cpu_ca(ca::CA.ComponentArray)
+
+Move ComponentArray form gpu to cpu.    
+"""
+function cpu_ca end
+# define in FluxExt
+
+

test/runtests.jl

@@ -3,6 +3,7 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml
 
 @time begin
     if GROUP == "All" || GROUP == "Basic"
+        @time @safetestset "test_HybridProblem" include("test_HybridProblem.jl")
         #@safetestset "test" include("test/test_ComponentArrayInterpreter.jl")
         @time @safetestset "test_ComponentArrayInterpreter" include("test_ComponentArrayInterpreter.jl")
         #@safetestset "test" include("test/test_gencovar.jl")

test/test_ComponentArrayInterpreter.jl

@@ -7,13 +7,15 @@ using ComponentArrays: ComponentArrays as CA
     component_counts = comp_cnts = (; P=2, M=3, Unc=5)
     m = ComponentArrayInterpreter(; comp_cnts...)
     testm = (m) -> begin
-        @test CM._get_ComponentArrayInterpreter_axes(m) == (CA.Axis(P=1:2, M=3:5, Unc=6:10),)
+        #type of axes may differ
+        #@test CM._get_ComponentArrayInterpreter_axes(m) == (CA.Axis(P=1:2, M=3:5, Unc=6:10),)
         @test length(m) == 10
         v = 1:length(m)
         cv = m(v)
         @test cv.Unc == 6:10
     end
     testm(m)
+    m = get_concrete(m)
     testm(get_concrete(m))
     Base.isconcretetype(typeof(m))
 end;

test/test_Flux.jl

@@ -1,6 +1,7 @@
 using Test
 using StatsFuns: logistic
 using CUDA, GPUArraysCore
+using ComponentArrays: ComponentArrays as CA
 
 using HybridVariationalInference
 # @testset "get_default_GPUHandler before loading Flux" begin
@@ -53,3 +54,10 @@ end;
     @test size(y) == (n_out, n_site)
 end;
 
+@testset "cpu_ca" begin
+    c1 = CA.ComponentVector(a=(a1=1,a2=2:3),b=3:4)
+    c1_gpu = gpu(c1)
+    #cpu(c1_gpu) # fails
+    @test cpu_ca(c1_gpu) == c1
+end;
+