Skip to content

Commit bf58655

Browse files
bgctwbgctw
committed
benchmark several alternatives in sampling
1 parent 267fc1c commit bf58655

9 files changed

Lines changed: 146 additions & 96 deletions

‎src/HVIApproximation.jl‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,13 @@ abstract type AbstractHVIApproximation end
88

99
abstract type AbstractMeanHVIApproximation <: AbstractHVIApproximation end
1010

11-
struct MeanHVIApproximation <: AbstractMeanHVIApproximation end
11+
# First implementation with one big sparse covariance matrix
1212
struct MeanHVIApproximationMat <: AbstractMeanHVIApproximation end
13-
struct MeanHVIApproximationDev <: AbstractMeanHVIApproximation end
13+
14+
# Reimplementation with generating random numbers for each block separately
15+
struct MeanHVIApproximation <: AbstractMeanHVIApproximation end
16+
17+
# for benchmarking changes, bevore implementing them
18+
struct MeanHVIApproximationDev <: AbstractMeanHVIApproximation end
1419

1520

‎src/HybridProblem.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ struct HybridProblem <: AbstractHybridProblem
5858
n_batch::Int,
5959
cor_ends::NamedTuple = (P = [length(ϕq[Val(:μP)])], M = [length(θM)]),
6060
pbm_covars::NTuple{N,Symbol} = (),
61-
approx::AbstractHVIApproximation = MeanHVIApproximation()
61+
approx::AbstractHVIApproximation = MeanHVIApproximationMat()
6262
) where N
6363
new(
6464
θM, f_batch, g, ϕg, ϕq, priors, py, transM, transP, cor_ends,
@@ -119,7 +119,7 @@ function update_hybridProblem(prob::AbstractHybridProblem; scenario,
119119
ϕq = get_hybridproblem_ϕq(prob; scenario),
120120
θP = nothing,
121121
ϕunc = nothing,
122-
approx::AbstractHVIApproximation = MeanHVIApproximation(),
122+
approx::AbstractHVIApproximation = MeanHVIApproximationMat(),
123123
)
124124
cor_ends_new = if !isnothing(cor_ends)
125125
# if new cor_ends was specified then re-initialize the ρsP and ρsM in ϕq

‎src/HybridVariationalInference.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,9 +110,9 @@ export get_ca_starts, get_ca_ends, get_cor_count
110110
include("cholesky.jl")
111111

112112
export neg_elbo_gtf, sample_posterior, predict_hvi, zero_penalty_loss
113+
include("elbo_dev.jl")
113114
include("elbo.jl")
114115
include("elbo2.jl")
115-
include("elbo_dev.jl")
116116

117117
export init_hybrid_params, init_hybrid_ϕunc
118118
include("init_hybrid_params.jl")

‎src/elbo.jl‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -472,31 +472,29 @@ function generate_ζ(approx::AbstractMeanHVIApproximation, rng::AbstractRNG,
472472
xMP0 = _append_each_covars(xM, CA.getdata(μ_ζP), pbm_covar_indices)
473473
μ_ζMs0 = g(xMP0, ϕg; is_testmode)
474474
ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(approx, rng, μ_ζP, μ_ζMs0, ϕq; n_MC, cor_ends, int_ϕq)
475+
ζsP = isempty(μ_ζP) ? ζP_resids : (μ_ζP .+ ζP_resids) # n_par x n_MC
475476
if pbm_covar_indices isa SA.SVector{0}
476477
# do not need to predict again but just add the residuals to μ_ζP and μ_ζMs
477478
#ζsP = μ_ζP .+ ζP_resids # n_par x n_MC # .+ on empty view does not work
478-
ζsP = isempty(μ_ζP) ? ζP_resids : (μ_ζP .+ ζP_resids) # n_par x n_MC
479479
ζsMs = permutedims(μ_ζMs0 .+ ζMs_parfirst_resids, (2, 1, 3)) # n_site x n_par x n_MC
480-
if any(ζsMs[:,2,:] .> 80.0)
481-
@show ζsMs
482-
@show ζMs_parfirst_resids
483-
@show ϕc.ϕq.coef_logσ2_ζMs
484-
error("encountered scaled residual outside envisoned range. Debug")
485-
end
480+
# if any(ζsMs[:,2,:] .> 80.0)
481+
# @show ζsMs
482+
# @show ζMs_parfirst_resids
483+
# @show ϕc.ϕq.coef_logσ2_ζMs
484+
# error("encountered scaled residual outside envisoned range. Debug")
485+
# end
486486
else
487487
#rP, rMs = first(zip(eachcol(ζP_resids), eachslice(ζMs_parfirst_resids;dims=3)))
488-
ζst = map(eachcol(ζP_resids), eachslice(ζMs_parfirst_resids; dims=3)) do rP, rMs
489-
ζP = μ_ζP .+ rP
488+
ζsMs_vec = map(eachcol(ζsP), eachslice(ζMs_parfirst_resids; dims=3)) do ζP, rMs
490489
# second pass: append ζP rather than μ_ζP to covars to xM
491490
xMP = _append_each_covars(xM, CA.getdata(ζP), pbm_covar_indices)
492491
μ_ζMst = g(xMP, ϕg)
493492
ζMs = (μ_ζMst .+ rMs)' # already transform to par-last form
494-
ζP, ζMs
493+
ζMs
495494
end
496495
# ζsP = stack(map(first, ζst); dims=1) # n_MC x n_par
497496
# ζsMs = stack(map(x -> x[2], ζst); dims=1) # n_MC x n_site x n_par
498-
ζsP = stack(map(first, ζst)) # n_par x n_MC
499-
ζsMs = stack(map(x -> x[2], ζst)) # n_site x n_par x n_MC
497+
ζsMs = stack(ζsMs_vec) # n_site x n_par x n_MC
500498
end
501499
ζsP, ζsMs, σ
502500
end

‎src/elbo2.jl‎

Lines changed: 47 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,28 @@
1+
# Re-Implementation of MeanHVIApproximationVec where random numbers are
2+
# generated for each block in the correlation matrix separately.
3+
#
4+
# Unfortunately, having many smaller matrix multiplications, the derivative
5+
# computation is slower than with
6+
# one big matrix multiplication of the big sparse blocked covariance matrix.
7+
# As benchmarked 251121 for 20 sites in test_sample_zeta.
8+
# Both on cpu and gpu
9+
#
10+
# For a large number of sites in a forward run across all sites
11+
# this implementaion is faster on gpu, but slower than the Vec
12+
# implementation on CPU (that uses triangular and blocked matrices)
13+
#
14+
# Hence, currently, rather use the MeanHVIApproximationVec CPU implementation.
15+
# for gradient on few sites (1.6ms cpu vs 5.4ms gpu)
16+
# and forward runs for many sites (760mus cpu vs 439ms gpu unit!)
17+
118
function sample_ζresid_norm(app::MeanHVIApproximation,
219
zP::AbstractMatrix, zMs::AbstractMatrix, ζP::TP, ζMs::TM,
320
ϕq::AbstractVector;
421
int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
522
cor_ends
623
) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}
7-
ϕuncc = int_ϕq(CA.getdata(ϕq))
24+
ϕunc = CA.getdata(ϕq)
25+
ϕuncc = int_ϕq(ϕunc)
826
n_θP, n_θMs, (n_θM, n_batch), n_MC = length(ζP), length(ζMs), size(ζMs), size(zP,1)
927
# do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
1028
ρsP = isempty(ϕuncc[Val(:ρsP)]) ? similar(ϕuncc[Val(:ρsP)]) : ϕuncc[Val(:ρsP)] # required by zygote
@@ -21,26 +39,46 @@ function sample_ζresid_norm(app::MeanHVIApproximation,
2139
σP = exp.(logσ2_ζP ./ 2)
2240
# create random numbers from U diag(σ) z = U (σ .* z)
2341
# for each block separately
24-
#σzMs = σMs .* reshape(zMs, (n_θM, n_batch, n_MC))
42+
#Ui, ri = first(zip(UPs, rangesP))
43+
ζP1 = ChainRulesCore.@ignore_derivatives similar(CA.getdata(ϕunc), size(zP'))
44+
#@benchmark typeof(zP .* ϕq[1:1]) # more allocations and more time
45+
local ζP_resids = if isempty(ζP)
46+
ζP1 # provide init of correct empty matrix type
47+
else
48+
mapreduce(vcat, UPs, rangesP) do Ui, ri
49+
Diagonal(σP[ri]) * Ui' * zP[:,ri]'
50+
end
51+
end::typeof(ζP1)
52+
#
2553
#Ui, ri = first(zip(UMs, rangesM))
2654
zMs_subjects = reshape(zMs, (n_MC, n_θM, n_batch))
2755
#σM, zM = first(zip(eachcol(σMs), eachslice(zMs_subjects; dims=3)))
2856
cat3 = (x,y) -> cat(x,y,dims=3)
2957
# map across subjects (n_batch)
3058
#ζMs_vec = map(eachcol(σMs), eachslice(zMs_subjects; dims=3)) do σM, zM
59+
ζMs1 = ChainRulesCore.@ignore_derivatives similar(ϕunc, size(zMs_subjects[:,:,1]'))
3160
fBlock = let UMs = UMs, rangesM = collect(rangesM) # without collect, type unstable
32-
function fBlock_inner(σM, zM)
61+
function fBlock_inner(σM, zM)::typeof(ζMs1)
62+
# generator does not work with Zygote
63+
# vcat(...) is not faster than mapreduce
64+
# g = map( UMs, rangesM) do Ui, ri # n_θM, n_MC
65+
# #(zM[:,ri]' * Ui * diagm(σM[ri]))'
66+
# diagm(σM[ri]) * Ui' * zM[:, ri]'
67+
# end
68+
# vcat(g...)
3369
mapreduce(vcat, UMs, rangesM) do Ui, ri # n_θM, n_MC
34-
ri::UnitRange{Int}
35-
#(zM[:,ri]' * Ui * diagm(σM[ri]))'
36-
diagm(σM[ri]) * Ui' * zM[:, ri]'
70+
# (zM[:,ri]' * (Ui * Diagonal(σM[ri])))'
71+
# Diagonal is faster than diagm and works with GPU
72+
Diagonal(σM[ri]) * Ui' * zM[:, ri]' # transpose to not return adjoint
3773
end
3874
end
3975
end
40-
ζMs1 = ChainRulesCore.@ignore_derivatives fBlock(σMs[:,1], zMs_subjects[:,:,1])
76+
#ζMs1 = ChainRulesCore.@ignore_derivatives fBlock(σMs[:,1], zMs_subjects[:,:,1])
77+
#TMR = ChainRulesCore.@ignore_derivatives typeof(similar(ϕq, size(zMs)))
4178
#TB = Base.infer_return_type(fBlock, (typeof(σMs[:,1]), typeof(zMs_subjects[:,:,1])))
4279
#ζMs_vec = map(fBlock, eachcol(σMs[:,2:end]), eachslice(zMs_subjects[:,:,2:end]; dims=3), init = ζMs1)
43-
ζMs_vec = map((σM, zM) -> fBlock(σM, zM)::typeof(ζMs1), eachcol(σMs), eachslice(zMs_subjects; dims=3))
80+
#ζMs_vec = map((σM, zM) -> fBlock(σM, zM)::TMR, eachcol(σMs), eachslice(zMs_subjects; dims=3))
81+
ζMs_vec = map((σM, zM) -> fBlock(σM, zM), eachcol(σMs), eachslice(zMs_subjects; dims=3))
4482
#zM = zMs_subjects[:,:,1]
4583
#ζMs_vec = [fBlock(σMs[:,i], zM) for i in axes(σMs, 2)]
4684
#ζMs_vec = [fBlock(σMs[:,1], zM)]
@@ -62,25 +100,6 @@ function sample_ζresid_norm(app::MeanHVIApproximation,
62100
# end
63101
# ζMs_parfirst_resids = reshape(ζMs_resids_stacked, n_θM, n_batch, n_MC)
64102
#
65-
#σzP = σP .* reshape(zP, (n_θP, n_MC))
66-
#Ui, ri = first(zip(UPs, rangesP))
67-
local ζP_resids = if isempty(ζP)
68-
# provide init of correct empty matrix type
69-
# single branch with init to mapreduce has Problems with Zygote
70-
ChainRulesCore.@ignore_derivatives ζMs_parfirst_resids[1:0,1,1:n_MC]
71-
else
72-
mapreduce(vcat, UPs, rangesP) do Ui, ri
73-
diagm(σP[ri]) * Ui' * zP[:,ri]'
74-
end
75-
end
76-
#
77-
() -> begin
78-
tmp = tmp1 = (σzP[ri,:]' * Ui)'
79-
tmp = tmp2 = (zP * (Ui * diagm(σP)))'
80-
tmp = tmp3 = (zP * Ui * diagm(σP))' # σ on wrong side of U
81-
tmp = tmp4 = diagm(σP) * Ui' * zP' # σ on wrong side of U
82-
std(tmp[2,:])
83-
end
84103
diagUσ = vcat(σP, vec(σMs))
85104
ζP_resids, ζMs_parfirst_resids, diagUσ
86-
end
105+
end

‎src/elbo_dev.jl‎

Lines changed: 30 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
1+
12
function sample_ζresid_norm(app::MeanHVIApproximationDev,
23
zP::AbstractMatrix, zMs::AbstractMatrix, ζP::TP, ζMs::TM,
34
ϕq::AbstractVector;
45
int_ϕq=get_concrete(ComponentArrayInterpreter(ϕq)),
56
cor_ends
67
) where {T,TP<:AbstractVector{T},TM<:AbstractMatrix{T}}
7-
ϕuncc = int_ϕq(CA.getdata(ϕq))
8+
ϕunc = CA.getdata(ϕq)
9+
ϕuncc = int_ϕq(ϕunc)
810
n_θP, n_θMs, (n_θM, n_batch), n_MC = length(ζP), length(ζMs), size(ζMs), size(zP,1)
911
# do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
1012
ρsP = isempty(ϕuncc[Val(:ρsP)]) ? similar(ϕuncc[Val(:ρsP)]) : ϕuncc[Val(:ρsP)] # required by zygote
@@ -21,26 +23,46 @@ function sample_ζresid_norm(app::MeanHVIApproximationDev,
2123
σP = exp.(logσ2_ζP ./ 2)
2224
# create random numbers from U diag(σ) z = U (σ .* z)
2325
# for each block separately
24-
#σzMs = σMs .* reshape(zMs, (n_θM, n_batch, n_MC))
26+
#Ui, ri = first(zip(UPs, rangesP))
27+
ζP1 = ChainRulesCore.@ignore_derivatives similar(CA.getdata(ϕunc), size(zP'))
28+
#@benchmark typeof(zP .* ϕq[1:1]) # more allocations and more time
29+
local ζP_resids = if isempty(ζP)
30+
ζP1 # provide init of correct empty matrix type
31+
else
32+
mapreduce(vcat, UPs, rangesP) do Ui, ri
33+
Diagonal(σP[ri]) * Ui' * zP[:,ri]'
34+
end
35+
end::typeof(ζP1)
36+
#
2537
#Ui, ri = first(zip(UMs, rangesM))
2638
zMs_subjects = reshape(zMs, (n_MC, n_θM, n_batch))
2739
#σM, zM = first(zip(eachcol(σMs), eachslice(zMs_subjects; dims=3)))
2840
cat3 = (x,y) -> cat(x,y,dims=3)
2941
# map across subjects (n_batch)
3042
#ζMs_vec = map(eachcol(σMs), eachslice(zMs_subjects; dims=3)) do σM, zM
43+
ζMs1 = ChainRulesCore.@ignore_derivatives similar(ϕunc, size(zMs_subjects[:,:,1]'))
3144
fBlock = let UMs = UMs, rangesM = collect(rangesM) # without collect, type unstable
32-
function fBlock_inner(σM, zM)
45+
function fBlock_inner(σM, zM)::typeof(ζMs1)
46+
# generator does not work with Zygote
47+
# vcat(...) is not faster than mapreduce
48+
# g = map( UMs, rangesM) do Ui, ri # n_θM, n_MC
49+
# #(zM[:,ri]' * Ui * diagm(σM[ri]))'
50+
# diagm(σM[ri]) * Ui' * zM[:, ri]'
51+
# end
52+
# vcat(g...)
3353
mapreduce(vcat, UMs, rangesM) do Ui, ri # n_θM, n_MC
34-
ri::UnitRange{Int}
35-
#(zM[:,ri]' * Ui * diagm(σM[ri]))'
36-
diagm(σM[ri]) * Ui' * zM[:, ri]'
54+
# (zM[:,ri]' * (Ui * Diagonal(σM[ri])))'
55+
# Diagonal is faster than diagm and works with GPU
56+
Diagonal(σM[ri]) * Ui' * zM[:, ri]' # transpose to not return adjoint
3757
end
3858
end
3959
end
40-
ζMs1 = ChainRulesCore.@ignore_derivatives fBlock(σMs[:,1], zMs_subjects[:,:,1])
60+
#ζMs1 = ChainRulesCore.@ignore_derivatives fBlock(σMs[:,1], zMs_subjects[:,:,1])
61+
#TMR = ChainRulesCore.@ignore_derivatives typeof(similar(ϕq, size(zMs)))
4162
#TB = Base.infer_return_type(fBlock, (typeof(σMs[:,1]), typeof(zMs_subjects[:,:,1])))
4263
#ζMs_vec = map(fBlock, eachcol(σMs[:,2:end]), eachslice(zMs_subjects[:,:,2:end]; dims=3), init = ζMs1)
43-
ζMs_vec = map((σM, zM) -> fBlock(σM, zM)::typeof(ζMs1), eachcol(σMs), eachslice(zMs_subjects; dims=3))
64+
#ζMs_vec = map((σM, zM) -> fBlock(σM, zM)::TMR, eachcol(σMs), eachslice(zMs_subjects; dims=3))
65+
ζMs_vec = map((σM, zM) -> fBlock(σM, zM), eachcol(σMs), eachslice(zMs_subjects; dims=3))
4466
#zM = zMs_subjects[:,:,1]
4567
#ζMs_vec = [fBlock(σMs[:,i], zM) for i in axes(σMs, 2)]
4668
#ζMs_vec = [fBlock(σMs[:,1], zM)]
@@ -62,25 +84,6 @@ function sample_ζresid_norm(app::MeanHVIApproximationDev,
6284
# end
6385
# ζMs_parfirst_resids = reshape(ζMs_resids_stacked, n_θM, n_batch, n_MC)
6486
#
65-
#σzP = σP .* reshape(zP, (n_θP, n_MC))
66-
#Ui, ri = first(zip(UPs, rangesP))
67-
local ζP_resids = if isempty(ζP)
68-
# provide init of correct empty matrix type
69-
# single branch with init to mapreduce has Problems with Zygote
70-
ChainRulesCore.@ignore_derivatives ζMs_parfirst_resids[1:0,1,1:n_MC]
71-
else
72-
mapreduce(vcat, UPs, rangesP) do Ui, ri
73-
diagm(σP[ri]) * Ui' * zP[:,ri]'
74-
end
75-
end
76-
#
77-
() -> begin
78-
tmp = tmp1 = (σzP[ri,:]' * Ui)'
79-
tmp = tmp2 = (zP * (Ui * diagm(σP)))'
80-
tmp = tmp3 = (zP * Ui * diagm(σP))' # σ on wrong side of U
81-
tmp = tmp4 = diagm(σP) * Ui' * zP' # σ on wrong side of U
82-
std(tmp[2,:])
83-
end
8487
diagUσ = vcat(σP, vec(σMs))
8588
ζP_resids, ζMs_parfirst_resids, diagUσ
8689
end

‎test/test_HybridProblem.jl‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,10 @@ function construct_problem(; scenario::Val{scen}) where scen
9191
xPvec=xP[:,1])
9292
ϕunc0 = init_hybrid_ϕunc(cor_ends, zero(FT))
9393
ϕq = CP.update_μP_by_θP(ϕunc0, θP, transP)
94-
approx = if (:MeanHVIApproxMat ∈ scen)
95-
MeanHVIApproximationMat()
96-
else
94+
approx = if (:MeanHVIApproxBlocks ∈ scen)
9795
MeanHVIApproximation()
96+
else
97+
MeanHVIApproximationMat()
9898
end
9999
HybridProblem(θM, ϕq, g_chain_scaled, ϕg0,
100100
f_batch, priors_dict, py,
@@ -381,7 +381,7 @@ test_with_flux_gpu = (scenario) -> begin
381381
end # if gdev isa MLDataDevices.AbstractGPUDevice
382382
end # test_with flux
383383

384-
#test_with_flux_gpu(Val((:MeanHVIApproxMat,))) # do not test any more
384+
#test_with_flux_gpu(Val((:MeanHVIApproxBlocks,))) # do not test any more, its slower
385385
#scenario = Val(())
386386
test_with_flux_gpu(Val((:default,)))
387387
test_with_flux_gpu(Val((:covarK2,)))

‎test/test_elbo.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ const prob = DoubleMM.DoubleMMCase()
2828
scenario = Val((:default,))
2929
#scenario = Val((:covarK2,))
3030

31-
#const approx = MeanHVIApproximationMat()
32-
const approx = MeanHVIApproximation()
31+
const approx = MeanHVIApproximationMat()
32+
#const approx = MeanHVIApproximation()
3333

3434
test_scenario = (scenario) -> begin
3535
probc = HybridProblem(prob; scenario, approx);

0 commit comments

Comments
 (0)