keep cleaning

Author: zuhengxu

.gitignore

@@ -3,6 +3,7 @@
 /docs/build/
 test/Manifest.toml
 example/Manifest.toml
+example/LocalPreferences.toml
 
 # Files generated by invoking Julia with --code-coverage
 *.jl.cov

example/Project.toml

@@ -21,6 +21,7 @@ LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 PlotlyJS = "f0f68f2c-4968-5e81-91da-67840de0976a"
@@ -36,3 +37,6 @@ TickTock = "9ff05d80-102d-5586-aa04-3a8bd1a90d20"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
+
+[extras]
+CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"

example/common.jl

@@ -47,38 +47,38 @@ function compare_trained_and_untrained_flow(
     return p
 end
 
-function check_trained_flow(
-    flow_trained::Bijectors.MultivariateTransformed,
-    true_dist::ContinuousMultivariateDistribution,
-    n_samples::Int;
-    kwargs...,
-)
-    samples_trained = rand_batch(flow_trained, n_samples)
-    samples_true = rand(true_dist, n_samples)
-
-    p = Plots.scatter(
-        samples_true[1, :],
-        samples_true[2, :];
-        label="True Distribution",
-        color=:green,
-        markersize=2,
-        alpha=0.5,
-    )
-    Plots.scatter!(
-        p,
-        samples_trained[1, :],
-        samples_trained[2, :];
-        label="Trained Flow",
-        color=:red,
-        markersize=2,
-        alpha=0.5,
-    )
-    Plots.plot!(; kwargs...)
-
-    Plots.title!(p, "Trained HamFlow")
-
-    return p
-end
+# function check_trained_flow(
+#     flow_trained::Bijectors.MultivariateTransformed,
+#     true_dist::ContinuousMultivariateDistribution,
+#     n_samples::Int;
+#     kwargs...,
+# )
+#     samples_trained = rand_batch(flow_trained, n_samples)
+#     samples_true = rand(true_dist, n_samples)
+
+#     p = Plots.scatter(
+#         samples_true[1, :],
+#         samples_true[2, :];
+#         label="True Distribution",
+#         color=:green,
+#         markersize=2,
+#         alpha=0.5,
+#     )
+#     Plots.scatter!(
+#         p,
+#         samples_trained[1, :],
+#         samples_trained[2, :];
+#         label="Trained Flow",
+#         color=:red,
+#         markersize=2,
+#         alpha=0.5,
+#     )
+#     Plots.plot!(; kwargs...)
+
+#     Plots.title!(p, "Trained HamFlow")
+
+#     return p
+# end
 
 function create_flow(Ls, q₀)
     ts = fchain(Ls)
@@ -89,33 +89,33 @@ end
 # training function for InvertibleNetworks
 ########################
 
-function pm_next!(pm, stats::NamedTuple)
-    return ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stats)])
-end
+# function pm_next!(pm, stats::NamedTuple)
+#     return ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stats)])
+# 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
-    )
+# 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 = []
+#     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
+#     # 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)
+#         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)
+#         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
+#         stat = (iteration=i, neg_log_llh=ls, gradient_norm=grad_norm)
+#         pm_next!(prog, stat)
+#     end
+#     return nnls
+# end

example/planar_radial_flow/planar_flow_main.jl

@@ -1,20 +1,22 @@
 using Random, Distributions, LinearAlgebra, Bijectors
 using ADTypes
 using Optimisers
-using FunctionChains
 using NormalizingFlows
-using Zygote
+using Mooncake
+using CUDA
 using Flux: f32
+using Flux
 using Plots
-include("../common.jl")
+include("common.jl")
 
 Random.seed!(123)
 rng = Random.default_rng()
+T = Float32
 
 ######################################
 # 2d Banana as the target distribution
 ######################################
-include("../targets/banana.jl")
+include("targets/banana.jl")
 
 # create target p
 p = Banana(2, 1.0f-1, 100.0f0)
@@ -26,13 +28,15 @@ logp = Base.Fix1(logpdf, p)
 function create_planar_flow(n_layers::Int, q₀)
     d = length(q₀)
     Ls = [f32(PlanarLayer(d)) for _ in 1:n_layers]
-    ts = fchain(Ls)
+    ts = reduce(∘, Ls)
     return transformed(q₀, ts)
 end
 
 # create a 10-layer planar flow
-flow = create_planar_flow(20, MvNormal(zeros(Float32, 2), I))
-flow_untrained = deepcopy(flow)
+q0 = MvNormal(zeros(T, 2), ones(T, 2))
+flow = create_planar_flow(10, q0)
+
+
 
 # train the flow
 sample_per_iter = 10

example/targets/banana.jl

@@ -2,39 +2,6 @@ using Distributions, Random
 using Plots
 using IrrationalConstants
 
-"""
-    Banana{T<:Real}
-
-Multidimensional banana-shape distribution.
-
-# Fields 
-$(FIELDS)
-
-# Explanation
-
-The banana distribution is obtained by applying a transformation ϕ to a multivariate normal 
-distribution ``\\mathcal{N}(0, \\text{diag}(var, 1, 1, …, 1))``. The transformation ϕ is defined as
-```math
-\phi(x_1, … , x_p) = (x_1, x_2 - B x_1^² + \text{var}*B, x_3, … , x_p)
-````
-which has a unit Jacobian determinant.
-
-Hence the density "fb" of a p-dimensional banana distribution is given by
-```math
-fb(x_1, \dots, x_p) = \exp\left[ -\frac{1}{2}\frac{x_1^2}{\text{var}} -
-\frac{1}{2}(x_2 + B x_1^2 - \text{var}*B)^2 - \frac{1}{2}(x_3^2 + x_4^2 + \dots
-+ x_p^2) \right] / Z,
-```
-where "B" is the "banananicity" constant, determining the curvature of a banana, and   
-``Z = \\sqrt{\\text{var} * (2\\pi)^p)}`` is the normalization constant.
-
-
-# Reference
-
-Gareth O. Roberts and Jeffrey S. Rosenthal
-"Examples of Adaptive MCMC."
-Journal of computational and graphical statistics, Volume 18, Number 2 (2009): 349-367.
-"""
 struct Banana{T<:Real} <: ContinuousMultivariateDistribution
     "Dimension of the distribution, must be >= 2"
     dim::Int      # Dimension