restructure example folder and refactor planar and radial examples

Author: zuhengxu

example/Project.toml

@@ -11,6 +11,7 @@ ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 FunctionChains = "8e6b2b91-af83-483e-ba35-d00930e4cf9b"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -19,6 +20,7 @@ IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"

example/SyntheticTargets.jl

@@ -1,5 +1,5 @@
+using DocStringExtensions
 using Distributions, Random, LinearAlgebra
-using LogDensityProblems
 using IrrationalConstants
 using Plots
 
@@ -23,7 +23,3 @@ function load_model(name::String)
         error("Model not defined")
     end
 end
-
-LogDensityProblems.dimension(dist::ContinuousDistribution) = length(dist)
-LogDensityProblems.logdensity(dist::ContinuousDistribution, x) = logpdf(dist, x)
-

example/common.jl

@@ -1,5 +1,4 @@
 using Random, Distributions, LinearAlgebra, Bijectors
-include("util.jl")
 
 function compare_trained_and_untrained_flow(
     flow_trained::Bijectors.MultivariateTransformed,
@@ -80,42 +79,16 @@ end
 #     return p
 # end
 
-function create_flow(Ls, q₀)
-    ts = fchain(Ls)
-    return transformed(q₀, ts)
-end
-
-#######################
-# training function for InvertibleNetworks
-########################
-
-# function pm_next!(pm, stats::NamedTuple)
-#     return ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stats)])
+# function create_flow(Ls, q₀)
+#     ts = fchain(Ls)
+#     return transformed(q₀, ts)
 # end
 
-# function train_invertible_networks!(G, loss, data_loader, n_epoch, opt)
-#     max_iters = n_epoch * length(data_loader)
-#     prog = ProgressMeter.Progress(
-#         max_iters; desc="Training", barlen=31, showspeed=true, enabled=true
-#     )
-
-#     nnls = []
-
-#     # training loop
-#     time_elapsed = @elapsed for (i, xs) in enumerate(IterTools.ncycle(data_loader, n_epoch))
-#         ls = loss(G, xs) #sets gradients of G
-
-#         push!(nnls, ls)
-
-#         grad_norm = 0
-#         for p in get_params(G)
-#             grad_norm += sum(abs2, p.grad)
-#             Flux.update!(opt, p.data, p.grad)
-#         end
-#         grad_norm = sqrt(grad_norm)
-
-#         stat = (iteration=i, neg_log_llh=ls, gradient_norm=grad_norm)
-#         pm_next!(prog, stat)
-#     end
-#     return nnls
-# end
+function visualize(p::Bijectors.MultivariateTransformed, samples=rand(p, 1000))
+    xrange = range(minimum(samples[1, :]) - 1, maximum(samples[1, :]) + 1; length=100)
+    yrange = range(minimum(samples[2, :]) - 1, maximum(samples[2, :]) + 1; length=100)
+    z = [exp(Distributions.logpdf(p, [x, y])) for x in xrange, y in yrange]
+    fig = contour(xrange, yrange, z'; levels=15, color=:viridis, label="PDF", linewidth=2)
+    scatter!(samples[1, :], samples[2, :]; label="Samples", alpha=0.3, legend=:bottomright)
+    return fig
+end

example/util.jl renamed to example/nn.jl

@@ -1,12 +1,4 @@
-using LinearAlgebra
-using StatsBase
-using BlockBandedMatrices
 using Flux, Bijectors
-using Base.Threads
-using DiffResults
-using ForwardDiff
-using ProgressMeter
-using TickTock
 
 function MLP_3layer(input_dim::Int, hdims::Int, output_dim::Int; activation=Flux.leakyrelu)
     return Chain(
@@ -55,12 +47,3 @@ end
 #     mlp_layer = MLP_BN(inputdim, hdim, outputdim; activation=activation)
 #     return Flux.SkipConnection(mlp_layer, +)
 # end
-
-function rand_batch(rng::AbstractRNG, td::Bijectors.MvTransformed, num_samples::Int)
-    samples = rand(rng, td.dist, num_samples)
-    res = td.transform(samples)
-    return res
-end
-function rand_batch(td::Bijectors.MvTransformed, num_samples::Int)
-    return rand_batch(Random.default_rng(), td, num_samples)
-end

example/planar_radial_flow/planar_flow_main.jl renamed to example/planar_flow_demo.jl

@@ -1,62 +1,64 @@
 using Random, Distributions, LinearAlgebra, Bijectors
-using ADTypes
-using Optimisers
+using Functors
+using Optimisers, ADTypes, Mooncake
 using NormalizingFlows
-using Mooncake
-using CUDA
-using Flux: f32
-using Flux
-using Plots
 include("common.jl")
+include("SyntheticTargets.jl")
 
 Random.seed!(123)
 rng = Random.default_rng()
-T = Float32
+T = Float64
 
 ######################################
 # 2d Banana as the target distribution
 ######################################
-include("targets/banana.jl")
+target = load_model("Banana")
+logp = Base.Fix1(logpdf, target)
 
-# create target p
-p = Banana(2, 1.0f-1, 100.0f0)
-logp = Base.Fix1(logpdf, p)
 
 ######################################
-# learn the target using planar flow 
+# setup planar flow
 ######################################
 function create_planar_flow(n_layers::Int, q₀)
     d = length(q₀)
-    Ls = [f32(PlanarLayer(d)) for _ in 1:n_layers]
+    Ls = [PlanarLayer(d) for _ in 1:n_layers]
     ts = reduce(∘, Ls)
     return transformed(q₀, ts)
 end
 
-# create a 10-layer planar flow
+@leaf MvNormal
 q0 = MvNormal(zeros(T, 2), ones(T, 2))
 flow = create_planar_flow(10, q0)
+flow_untrained = deepcopy(flow)
 
 
 
-# train the flow
-sample_per_iter = 10
-cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,)
-checkconv(iter, stat, re, θ, st) = stat.gradient_norm < 1e-3
+######################################
+# start training
+######################################
+sample_per_iter = 30
+
+# callback function to log training progress
+cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
+adtype = ADTypes.AutoMooncake(; config = Mooncake.Config())
+checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
 flow_trained, stats, _ = train_flow(
     elbo,
     flow,
     logp,
     sample_per_iter;
-    max_iters=200_00,
-    optimiser=Optimisers.Adam(),
+    max_iters=10_000,
+    optimiser=Optimisers.Adam(one(T)/100),
+    ADbackend=adtype,
+    show_progress=true,
     callback=cb,
-    ADbackend=AutoZygote(),
     hasconverged=checkconv,
 )
+θ, re = Optimisers.destructure(flow_trained)
 losses = map(x -> x.loss, stats)
 
 ######################################
 # evaluate trained flow
 ######################################
 plot(losses; label="Loss", linewidth=2) # plot the loss
-compare_trained_and_untrained_flow(flow_trained, flow_untrained, p, 1000)
+compare_trained_and_untrained_flow(flow_trained, flow_untrained, target, 1000)

example/planar_radial_flow/radial_flow_main.jl

@@ -0,0 +1,63 @@
+using Random, Distributions, LinearAlgebra, Bijectors
+using Functors
+using Optimisers, ADTypes, Mooncake
+using NormalizingFlows
+include("common.jl")
+include("SyntheticTargets.jl")
+
+Random.seed!(123)
+rng = Random.default_rng()
+T = Float64
+
+######################################
+# get target logp
+######################################
+target = load_model("WarpedGaussian")
+logp = Base.Fix1(logpdf, target)
+
+######################################
+# setup radial flow 
+######################################
+function create_radial_flow(n_layers::Int, q₀)
+    d = length(q₀)
+    Ls = [RadialLayer(d) for _ in 1:n_layers]
+    ts = reduce(∘, Ls)
+    return transformed(q₀, ts)
+end
+
+# create a 10-layer radial flow
+@leaf MvNormal
+q0 = MvNormal(zeros(T, 2), ones(T, 2))
+flow = create_radial_flow(10, q0)
+
+flow_untrained = deepcopy(flow)
+
+######################################
+# start training
+######################################
+sample_per_iter = 30
+
+# callback function to log training progress
+cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
+adtype = ADTypes.AutoMooncake(; config = Mooncake.Config())
+checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
+flow_trained, stats, _ = train_flow(
+    elbo,
+    flow,
+    logp,
+    sample_per_iter;
+    max_iters=10_000,
+    optimiser=Optimisers.Adam(one(T)/100),
+    ADbackend=adtype,
+    show_progress=true,
+    callback=cb,
+    hasconverged=checkconv,
+)
+θ, re = Optimisers.destructure(flow_trained)
+losses = map(x -> x.loss, stats)
+
+######################################
+# evaluate trained flow
+######################################
+plot(losses; label="Loss", linewidth=2) # plot the loss
+compare_trained_and_untrained_flow(flow_trained, flow_untrained, target, 1000)