more NF examples #11

zuhengxu · 2023-07-13T05:45:02Z

as per issue #15

implement a set of famous NFs illustrating the usage of NormalizingFlows.jl + Bijectors.jl (a tentative list)
- planar and radial flow
- Real NVP
- Neural Spline flow
- Sparse Hamiltonian flow

This reverts commit 1c1c88a.

…mooncake

yebai · 2025-04-10T21:53:57Z

Could we add a script to test whether all the example models run, e.g. on buildkite GPU?

zuhengxu · 2025-04-10T22:02:25Z

Yes, let me clean up the Ham VI demo and then add the test script. Currently I can confirm that all the examples run on CPUs by doing

include("demo_..._fllow.jl")

zuhengxu · 2025-04-12T01:24:27Z

Could we add a script to test whether all the example models run, e.g. on buildkite GPU?

@yebai Currently all the examples are only tested on CPUs. Do you have examples of such script (for testing runs) that I can refer to? (I did do include(...jl) manually for all the demos though.)

In terms of GPU, shall we first merge this PR and do GPU examples/test (#49) on a separate PR? the purpose this PR is just to demonstrate how to define and customize flow layers (at least for those popular flows). I think that to properly take advantage of GPU, we might want some interface adjustment.

yebai · 2025-04-14T14:35:38Z

Currently all the examples are only tested on CPUs. Do you have examples of such a script (for testing runs) that I can refer to? (I did do include(...jl) manually for all the demos though.)

It is probably okay to create a new CI group called NF Examples and include all examples there. If running all examples is time-consuming, we could run only a few iterations and not wait for the convergence. This is mainly to ensure all examples can still run with new releases or PRs.

In terms of GPU, shall we first merge this PR and do GPU examples/test (#49) on a separate PR?

No problem -- it sounds good!

sunxd3

General, this is great effort!

Some minor comments, I didn't run the code to see the results.

In the future, probably all of these should be in a documentation page that automatically run and render results. But this PR is fine as it is.
(Worth it to add an issue as TODO?)

sunxd3 · 2025-04-15T08:57:00Z

Project.toml

@@ -8,7 +8,9 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"


Is there a reason for adding these dependencies?

sunxd3 · 2025-04-15T08:57:07Z

Project.toml

 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"


(also this)

sunxd3 · 2025-04-15T09:00:26Z

example/SyntheticTargets.jl

+function load_model(name::String)
+    if name == "Banana"
+        return Banana(2, 1.0, 10.0)
+    elseif name == "Cross"
+        return Cross()
+    elseif name == "Funnel"
+        return Funnel(2)
+    elseif name == "WarpedGaussian"
+        return WarpedGauss()
+    else
+        error("Model not defined")
+    end
+end


this is slightly ad hoc (it required a jump to know what this function will return at the call site), maybe use model constructor directly?

sunxd3 · 2025-04-15T09:02:47Z

example/demo_hamiltonian_flow.jl

+    logp_ρ = sum(logpdf(Normal(), ρ))
+    return logp_x + logp_ρ
+end
+∇logp = Base.Fix1(score, target)


score is only defined for banana, worth mention a bit, I think

sunxd3 · 2025-04-15T09:03:49Z

example/utils.jl

+# 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


is this required anymore?

sunxd3 · 2025-04-15T09:05:08Z

example/Project.toml

@@ -2,16 +2,22 @@
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"


this might not be required anymore?

sunxd3 · 2025-04-15T09:06:55Z

example/Project.toml

 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a"


SimpleUnPack can be replaced with destructing syntax (like (; a, b) = x) https://github.com/devmotion/SimpleUnPack.jl/blob/5943a14470a08f919f82809536ac63aed67bb28f/src/SimpleUnPack.jl#L22

zuhengxu and others added 12 commits July 10, 2023 23:28

add support for hasconverged

4f591af

fix test error

4e31807

rm example/Manifest.toml

99ac0c1

minor bug fix for trainig loop

e4f5efa

test new stopping criterion

b345e78

test convergent condition/ rm unready examples

f10512f

Merge branch 'TuringLang:main' into hasconverge

b93dbe8

rm julia test from CI

1c1c88a

Revert "rm julia test from CI"

5d2844f

This reverts commit 1c1c88a.

make autodiff pkgs as extension + require for bwd compat

3fdcb0e

debugging Ext

ddad59e

keep debugging

ef60ee1

zuhengxu self-assigned this Jul 13, 2023

zuhengxu added documentation Improvements or additions to documentation enhancement New feature or request labels Jul 13, 2023

sunxd3 and others added 7 commits July 13, 2023 09:20

Fix AD package extension loading issues

5a5deb0

Applying @devmotion's comment

25d4211

patch last commit

df5eddd

patch for julia 1.6

d44143b

loading dep pkgs from main pkg instead of functions for explicitness

b03e922

fixing test err

e9acf70

rm unready examples

c6bf68b

zuhengxu changed the title ~~more NF examples+ more proper synthetic examples + documentations~~ more NF examples Jul 13, 2023

zuhengxu added 7 commits July 23, 2023 20:58

update realnvp

fabf20a

minor ed

7964f91

removing unnecessary import

01de9d4

refactor affinecoupling and example/

6993d33

debug affinecoupling flow

810881f

adapt to the updated autoforwarddiff to resolve test err

0277c9e

fix test err

4bc02bf

zuhengxu added 12 commits March 16, 2025 14:06

update

b30d961

Merge branch 'stability' into more_examples

aeda365

rm many unrelated code

2c4d793

merge conflict from upstream

6056ef1

keep cleaning

960a359

merge from turing/main

51bc5b9

fix warpgaussian logpdf error/making neal funnel logpdf working with …

f5a2f9a

…mooncake

add easier model loading code

d916b64

restructure example folder and refactor planar and radial examples

d84fce4

clean demos for realnvp/planar/radial/ fix a bug in nsf

6872760

rm enzyme dependency

e3fb428

tune nsf flow--enlarge B--to make it work

ddf0ba9

zuhengxu marked this pull request as ready for review April 10, 2025 21:46

rm some useless files from HamVI

f2affe2

zuhengxu added 8 commits April 10, 2025 15:31

rename MLP_3layer to mlp3 for convenience

78d5f45

rename common.jl to utils

2229577

minor update file naming, and better nsf implementation

aefc651

minor ed

7f1580c

rm redundant nsf file

750c1e1

shrink stepsize for nsf

ab48d11

cleaned hamiltonian flow

42bc4a1

minor ed

edf1bac

zuhengxu requested a review from sunxd3 April 12, 2025 01:24

rm redundant pkgs

c05789b

yebai approved these changes Apr 14, 2025

View reviewed changes

sunxd3 requested changes Apr 15, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

more NF examples #11

more NF examples #11

zuhengxu commented Jul 13, 2023 •

edited

Loading

yebai commented Apr 10, 2025

zuhengxu commented Apr 10, 2025

zuhengxu commented Apr 12, 2025 •

edited

Loading

yebai commented Apr 14, 2025 •

edited

Loading

sunxd3 left a comment

sunxd3 Apr 15, 2025

sunxd3 Apr 15, 2025

sunxd3 Apr 15, 2025

sunxd3 Apr 15, 2025

sunxd3 Apr 15, 2025

sunxd3 Apr 15, 2025

sunxd3 Apr 15, 2025

		LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
		MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"

more NF examples #11

Are you sure you want to change the base?

more NF examples #11

Conversation

zuhengxu commented Jul 13, 2023 • edited Loading

yebai commented Apr 10, 2025

zuhengxu commented Apr 10, 2025

zuhengxu commented Apr 12, 2025 • edited Loading

yebai commented Apr 14, 2025 • edited Loading

sunxd3 left a comment

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

sunxd3 Apr 15, 2025

Choose a reason for hiding this comment

zuhengxu commented Jul 13, 2023 •

edited

Loading

zuhengxu commented Apr 12, 2025 •

edited

Loading

yebai commented Apr 14, 2025 •

edited

Loading