bug fix with CUDA and gradient calculations.

Author: RainerHeintzmann

examples/gauss_fit.jl

@@ -14,7 +14,8 @@ many_sig = true
 
 use_cuda = false
 DType = Float32
-N = 10_000
+# N = 10_000
+N = 1000
 hp_off = many_off ? 2 .*rand(DType, (1, N)) : 0
 hp_sig = many_sig ? zeros(DType, (1, N)) : 0
 hp_int = many_int ? 1 .+ rand(DType, (1, N)) : 1
@@ -57,7 +58,7 @@ myfg! = get_fg!(dat, gaussian_raw, length(sz); loss=loss_poisson_pos);
 #     sigma = CuArray(sigma)
 # end
 # startvals = DType.(ComponentVector(;bg=bg, intensity=intensity, off = soff, args = sigma))
-opt = Optim.Options(iterations = 1500); #
+opt = Optim.Options(iterations = 150); #
 
 if (false)
     G = copy(startvals)
@@ -86,6 +87,30 @@ reso.minimum #
 success = sum(abs.(collect(startvals.off) .- vec_true.off), dims=1) .< 0.5
 success = success .&& sum(abs.(collect(reso.minimizer.off) .- vec_true.off), dims=1) .< 0.5
 sum(.!success)
+
+# Try the same by calling optimizations in parallel (seems to work better!):
+if false
+    szz = size(dat)[end]
+    # convert the startvals into a vector of individual startvals
+    dat_a = Array(dat)
+    svv = [DType.(ComponentVector(gauss_start(dat_a[:,:,n:n], 0.2, length(sz))))  for n in 1:szz];
+    if (use_cuda)
+        svv = [ComponentVector(;bg=CuArray(svv[n].bg), intensity=CuArray(svv[n].intensity),
+            off = CuArray(svv[n].off), args = CuArray(svv[n].args)) for n in 1:szz]
+    end
+    # svb.args = svb.args .* 1.2f0
+    myfg!v = [get_fg!(dat[:,:,n:n], gaussian_raw, length(sz); loss=loss_poisson_pos) for n in 1:szz];
+
+    odov = [OnceDifferentiable(Optim.NLSolversBase.only_fg!(myfg!v[n]), svv[n]) for n in 1:szz];
+    @time resov = Optim.optimize.(odov, svv, Ref(Optim.LBFGS()), Ref(opt));
+    # seeems faster!
+    svvoff = cat([svv[n].off for n in 1:szz]..., dims=2)
+    resoff = cat([resov[n].minimizer.off for n in 1:szz]..., dims=2)
+    success = sum(abs.(collect(svvoff) .- vec_true.off), dims=1) .< 0.5
+    success = success .&& sum(abs.(collect(resoff) .- vec_true.off), dims=1) .< 0.5
+    # all worked!
+    sum(.!success)
+end
 # findfirst(.!success)
 @vt collect(dat)[:,:,.!success[:]] collect(pdat)[:,:,.!success[:]] collect(gaussian_vec(sz, startvals))[:,:,.!success[:]] collect(gaussian_vec(sz, reso.minimizer))[:,:,.!success[:]]
 

src/general.jl

@@ -58,21 +58,30 @@ function calculate_separables_nokw(::Type{AT}, fct, sz::NTuple{N, Int},
     # args_1d = ntuple((d) -> arg_n(d, args), Val(N))
     # in_place_assing!.(all_axes, 1, fct, idcs, sz, args_1d)
     for (res, sz1d, d) in zip(all_axes, sz, 1:N)
-        off = get_vec_dim(offset, d, sz)
-        sca = get_vec_dim(scale, d, sz)
+        # off = get_vec_dim(offset, d, sz) # not needed any more since in get_1d_ids
+        # sca = get_vec_dim(scale, d, sz)
         idc = get_1d_ids(d, sz, offset, scale)
         args_d = arg_n(d, args, RT, sz) # 
         # in_place_assing!(res, 1, fct, idc, sz1d, args_d)
         # @show size(res)
-        # @show size(idc)        
         res .= fct.(idc, sz1d, args_d...) # 5 allocs, 160 bytes
     end
     return all_axes
     # return res
 end
 
+"""
+    get_1d_ids(d, sz::NTuple{N, Int}, offset, scale) where {N}
+
+returns one-dimensional indices for a given dimension `d` of an N-dimensional array.
+The indices are shifted by `offset` and scaled by `scale`, which can also be vectors 
+"""
+# for Numbers, the reorient comes last, to have it CUDA-compatible
+get_1d_ids(d, sz::NTuple{N, Int}, offset::Number, scale::Number) where {N} = (reorient(get_vec_dim(scale, d, sz) .* ((1:sz[d]) .- get_vec_dim(offset, d, sz)), d, Val(N)))
+# for abstract arrays, we first have to reorient. 
 get_1d_ids(d, sz::NTuple{N, Int}, offset, scale) where {N} = get_vec_dim(scale, d, sz) .* (reorient((1:sz[d]), d, Val(N)) .- get_vec_dim(offset, d, sz))
-get_1d_ids(d, sz::NTuple{N, Int}, offset) where {N} = (reorient((1:sz[d]), d, Val(N)) .- get_vec_dim(offset, d, sz))
+get_1d_ids(d, sz::NTuple{N, Int}, offset::Number) where {N} = (reorient((1:sz[d]) .- get_vec_dim(offset, d, sz), d, Val(N)))
+get_1d_ids(d, sz::NTuple{N, Int}, offset) where {N} = reorient(1:sz[d], d, Val(N)) .- get_vec_dim(offset, d, sz)
 # get_1d_ids(d, sz, offset, scale) = pick_n(d, scale) .* ((1:sz[d]) .- pick_n(d, offset))
 # get_1d_ids(d, sz, offset::NTuple, scale::NTuple) = scale[d] .* ((1:sz[d]) .- offset[d])
 
@@ -340,6 +349,7 @@ function ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(cal
     sca = isnothing(args[2]) ? RAT([one(RT)]) : RT.(args[2])
 
     ids = ntuple((d) -> get_1d_ids(d, sz, off, sca), Val(N)) # offset==args[1] and scale==args[2]
+    # ids_offset_only = ntuple((d) -> get_1d_ids(d, sz, off), Val(N)) # offset==args[1] and scale==args[2]
     ids_offset_only = get_1d_ids.(1:N, Ref(sz), Ref(off)) # , one(eltype(AT))
 
     extra_sz = get_arg_sz(sz, args...)

test/speedtests.jl

@@ -3,6 +3,7 @@ using IndexFunArrays
 using SeparableFunctions
 using BenchmarkTools
 using CUDA
+using PlotlyJS
 
 function speedt_test()
     sz = (256, 256, 256)
@@ -15,7 +16,7 @@ function speedt_test()
         exp(-sum(abs2.((Tuple(ci) .- offset)./sigma)))
     end
     res = get_exp.(CartesianIndices(sz), Ref(Float32(sqrt(2)).*sigma), Ref(offset));  
-    @btime get_exp.(CartesianIndices($sz), Ref($sigma)); # 47.7 ms (2 allocations, 64 Mb) , but 243 ms with offset!
+    @btime get_exp.(CartesianIndices($sz), Ref($sigma), Ref(0)); # 47.7 ms (2 allocations, 64 Mb) , but 243 ms with offset!
     @btime get_exp.(CartesianIndices($sz), Ref($sigma), Ref(offset)); # 47.7 ms, but 243 ms with offset (7 allocations, 64 Mb)!
 
     res2 = similar(res);
@@ -62,7 +63,6 @@ function speedt_test()
     res3_sep = gaussian_sep(typeof(resc), sz, sigma=sigma, offset=offset); # 
     tc_res2_assign = @belapsed CUDA.@sync $resc .= $res3_sep
 
-    using PlotlyJS
     method = ["Compute In Place", "Collect Separables", "Lazy Arrays", "Collect Precomputed", "Precompute"]
     dat_no_cuda = 1000 .*[t_in_place, t_gaussian_col, t_gaussian_lz, t_res2_assign, t_gaussian_sep]
     dat_cuda = 1000 .*[tc_get_exp, tc_gaussian_col, tc_gaussian_lz, tc_res2_assign, tc_gaussian_sep]
@@ -73,7 +73,7 @@ function speedt_test()
     ], Layout(title="3D Gaussian (512x512x256)", yaxis=attr(title="Time [ms]", type="log"))) # barmode="stack", 
 
 
-    # now som speed comparison for propagator_col!
+    # now some speed comparison for propagator_col!
     sz = (1024, 1024)
     Δz = 1f0
     scale = 0.5f0 ./ (max.(sz ./ 2, 1))