Rational-quadratic spline transformations

.gitignore

@@ -5,3 +5,4 @@
 .ipynb_checkpoints
 .vscode
 Manifest.toml
+spline_testing.jl

Project.toml

@@ -13,6 +13,7 @@ ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"

notes.jl

@@ -0,0 +1,5 @@
+    # think about storage, so memory coalescence can be used?, also shared memory caching for older hardware
+    # keep data on gpu
+    # use profiling to really get in there with the optimization
+    #use NVIDIAs compute sanitizer to look at emory issues: cuda/compute-sanitizer
+    # instpection tools : @macroexpand use in front of @kernel in kernel definition; @ka_code_typed, @ka_code_llvm use in front of call to kernel

src/EuclidianNormalizingFlows.jl

@@ -29,6 +29,8 @@ using Parameters
 using SpecialFunctions
 using StatsBase
 using ValueShapes
+using KernelAbstractions
+using KernelAbstractions: @atomic
 
 import Zygote
 import ZygoteRules
@@ -46,5 +48,6 @@ include("householder_trafo.jl")
 include("scale_shift_trafo.jl")
 include("center_stretch.jl")
 include("johnson_trafo.jl")
+include("spline_trafo.jl")
 
 end # module

src/optimize_whitening.jl

@@ -6,6 +6,7 @@ std_normal_logpdf(x::Real) = -(abs2(x) + log2π)/2
 
 function mvnormal_negll_trafo(trafo::Function, X::AbstractMatrix{<:Real})
     nsamples = size(X, 2) # normalize by number of samples to be independent of batch size:
+
     Y, ladj = with_logabsdet_jacobian(trafo, X)
     #ref_ll = sum(sum(std_normal_logpdf.(Y), dims = 1) .+ ladj) / nsamples
     # Faster:
@@ -26,7 +27,8 @@ function optimize_whitening(
     smpls::VectorOfSimilarVectors{<:Real}, initial_trafo::Function, optimizer;
     nbatches::Integer = 100, nepochs::Integer = 100,
     optstate = Optimisers.setup(optimizer, deepcopy(initial_trafo)),
-    negll_history = Vector{Float64}()
+    negll_history = Vector{Float64}(),
+    shuffle_samples::Bool = false
 )
     batchsize = round(Int, length(smpls) / nbatches)
     batches = collect(Iterators.partition(smpls, batchsize))
@@ -40,6 +42,10 @@ function optimize_whitening(
             state, trafo = Optimisers.update(state, trafo, d_trafo)
             push!(negll_hist, negll)
         end
+        if shuffle_samples
+            shuffled_smpls = shuffle(smpls)
+            batches = collect(Iterators.partition(shuffled_smpls, batchsize))
+        end
     end
     (result = trafo, optimizer_state = state, negll_history = vcat(negll_history, negll_hist))
 end