'+ escapeHtml(title) + '
' + escapeHtml(summary) +'
diff --git a/.nojekyll b/.nojekyll new file mode 100644 index 0000000..e69de29 diff --git a/404.html b/404.html new file mode 100644 index 0000000..d00a850 --- /dev/null +++ b/404.html @@ -0,0 +1,143 @@ + + +
+ + + + + + + +Page not found
+normflows
: A PyTorch Package for Normalizing Flowsnormflows
is a PyTorch implementation of discrete normalizing flows. Many popular flow architectures are implemented,
+see the list below. The package can be easily installed via pip.
+The basic usage is described here, and a full documentation
+is available as well. A more detailed description of this package is given in out accompanying
+paper.
Several sample use cases are provided in the
+examples
folder,
+including Glow,
+a VAE, and
+a Residual Flow.
+Moreover, two simple applications are highlighed in the examples section. You can run them
+yourself in Google Colab using the links below to get a feeling for normflows
.
Architecture | +Reference | +
---|---|
Planar Flow | +Rezende & Mohamed, 2015 | +
Radial Flow | +Rezende & Mohamed, 2015 | +
NICE | +Dinh et al., 2014 | +
Real NVP | +Dinh et al., 2017 | +
Glow | +Kingma et al., 2018 | +
Masked Autoregressive Flow | +Papamakarios et al., 2017 | +
Neural Spline Flow | +Durkan et al., 2019 | +
Circular Neural Spline Flow | +Rezende et al., 2020 | +
Residual Flow | +Chen et al., 2019 | +
Stochastic Normalizing Flow | +Wu et al., 2020 | +
Note that Neural Spline Flows with circular and non-circular coordinates +are supported as well.
+The latest version of the package can be installed via pip
+pip install normflows
+
+At least Python 3.7 is required. If you want to use a GPU, make sure that +PyTorch is set up correctly by following the instructions at the +PyTorch website.
+To run the example notebooks clone the repository first
+git clone https://github.com/VincentStimper/normalizing-flows.git
+
+and then install the dependencies.
+pip install -r requirements_examples.txt
+
+A normalizing flow consists of a base distribution, defined in
+nf.distributions.base
,
+and a list of flows, given in
+nf.flows
.
+Let's assume our target is a 2D distribution. We pick a diagonal Gaussian
+base distribution, which is the most popular choice. Our flow shall be a
+Real NVP model and, therefore, we need
+to define a neural network for computing the parameters of the affine coupling
+map. One dimension is used to compute the scale and shift parameter for the
+other dimension. After each coupling layer we swap their roles.
import normflows as nf
+
+# Define 2D Gaussian base distribution
+base = nf.distributions.base.DiagGaussian(2)
+
+# Define list of flows
+num_layers = 32
+flows = []
+for i in range(num_layers):
+ # Neural network with two hidden layers having 64 units each
+ # Last layer is initialized by zeros making training more stable
+ param_map = nf.nets.MLP([1, 64, 64, 2], init_zeros=True)
+ # Add flow layer
+ flows.append(nf.flows.AffineCouplingBlock(param_map))
+ # Swap dimensions
+ flows.append(nf.flows.Permute(2, mode='swap'))
+
+Once they are set up, we can define a
+nf.NormalizingFlow
+model. If the target density is available, it can be added to the model
+to be used during training. Sample target distributions are given in
+nf.distributions.target
.
# If the target density is not given
+model = nf.NormalizingFlow(base, flows)
+
+# If the target density is given
+target = nf.distributions.target.TwoMoons()
+model = nf.NormalizingFlow(base, flows, target)
+
+The loss can be computed with the methods of the model and minimized.
+# When doing maximum likelihood learning, i.e. minimizing the forward KLD
+# with no target distribution given
+loss = model.forward_kld(x)
+
+# When minimizing the reverse KLD based on the given target distribution
+loss = model.reverse_kld(num_samples=512)
+
+# Optimization as usual
+loss.backward()
+optimizer.step()
+
+We provide several illustrative examples of how to use the package in the
+examples
+directory. Amoung them are implementations of
+Glow,
+a VAE, and
+a Residual Flow.
+More advanced experiments can be done with the scripts listed in the
+repository about resampled base distributions,
+see its experiments
+folder.
Below, we consider two simple 2D examples.
+In this notebook, +which can directly be opened in +Colab, +we consider a 2D distribution with two half-moon-shaped modes as a target. We approximate it with a Real NVP model +and obtain the following results.
+Note that there might be a density filament connecting the two modes, which is due to an architectural limitation +of normalizing flows, especially prominent in Real NVP. You can find out more about it in +this paper.
+In another example, +which is available in Colab +as well, we apply a Neural Spline Flow model to a distribution defined on a cylinder. The resulting density is visualized below.
+This example is considered in the paper accompanying this repository.
+If you have problems, please read the package documentation +and check out the examples section above. You are also welcome to +create issues on GitHub to get help. Note that it is +worthwhile browsing the existing open +and closed issues, which might +address the problem you are facing.
+If you find a bug or have a feature request, please +file an issue on GitHub.
+You are welcome to contribute to the package by fixing the bug or adding the feature yourself. If you want to
+contribute, please add tests for the code you added or modified and ensure it passes successfully by running pytest
.
+This can be done by simply executing
pytest
+
+within your local version of the repository. Make sure you code is well documented, and we also encourage contributions +to the existing documentation. Once you finished coding and testing, please +create a pull request on GitHub.
+The package has been used in several research papers, which are listed below.
+++Andrew Campbell, Wenlong Chen, Vincent Stimper, José Miguel Hernández-Lobato, and Yichuan Zhang. +A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization. +In Proceedings of the 38th International Conference on Machine Learning, pp. 1238–1248. PMLR, 2021.
+ +Vincent Stimper, Bernhard Schölkopf, José Miguel Hernández-Lobato. +Resampling Base Distributions of Normalizing Flows. +In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 4915–4936, 2022.
+ +Laurence I. Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, José Miguel Hernández-Lobato. +Flow Annealed Importance Sampling Bootstrap. +arXiv preprint arXiv:2208.01893, 2022.
+ +
Moreover, the boltzgen
package
+has been build upon normflows
.
If you use normflows
, please consider citing the corresponding paper as follows.
++Vincent Stimper, David Liu, Andrew Campbell, Vincent Berenz, Lukas Ryll, Bernhard Schölkopf, José Miguel Hernández-Lobato. +normflows: A PyTorch Package for Normalizing Flows, arXiv preprint arXiv:2302.12014, 2023.
+
Bibtex
+@article{normflows,
+ author = {Vincent Stimper and David Liu and Andrew Campbell and Vincent Berenz and Lukas Ryll and Bernhard Sch{\"o}lkopf and Jos{\'e} Miguel Hern{\'a}ndez-Lobato},
+ title = {normflows: {A} {P}y{T}orch {P}ackage for {N}ormalizing {F}lows},
+ journal = {arXiv preprint arXiv:2302.12014},
+ year = {2023}
+}
+
t |
216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 +292 +293 +294 +295 +296 +297 |
|
__init__(q0, flows)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
q0 |
+ + | +Base distribution |
+ + required + | +
flows |
+ + | +List of flows |
+ + required + | +
normflows/core.py
221 +222 +223 +224 +225 +226 +227 +228 +229 +230 |
|
forward_kld(x, y)
+
+Estimates forward KL divergence, see arXiv 1912.02762
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch sampled from target distribution |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Estimate of forward KL divergence averaged over batch |
+
normflows/core.py
232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 |
|
load(path)
+
+Load model from state dict
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
path |
+ + | +Path including filename where to load model from |
+ + required + | +
normflows/core.py
291 +292 +293 +294 +295 +296 +297 |
|
log_prob(x, y)
+
+Get log probability for batch
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch |
+ + required + | +
y |
+ + | +Classes of x |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability |
+
normflows/core.py
265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 |
|
sample(num_samples=1, y=None)
+
+Samples from flow-based approximate distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw |
+
+ 1
+ |
+
y |
+ + | +Classes to sample from, will be sampled uniformly if None |
+
+ None
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples, log probability |
+
normflows/core.py
249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 |
|
save(path)
+
+Save state dict of model
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
param |
+
+ path
+ |
+ Path including filename where to save model |
+ + required + | +
normflows/core.py
283 +284 +285 +286 +287 +288 +289 |
|
MultiscaleFlow
+
+
+
+ Bases: nn.Module
Normalizing Flow model with multiscale architecture, see RealNVP or Glow paper
+ + +normflows/core.py
300 +301 +302 +303 +304 +305 +306 +307 +308 +309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 +324 +325 +326 +327 +328 +329 +330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 +384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 |
|
__init__(q0, flows, merges, transform=None, class_cond=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
q0 |
+ + | +List of base distribution |
+ + required + | +
flows |
+ + | +List of list of flows for each level |
+ + required + | +
merges |
+ + | +List of merge/split operations (forward pass must do merge) |
+ + required + | +
transform |
+ + | +Initial transformation of inputs |
+
+ None
+ |
+
class_cond |
+ + | +Flag, indicated whether model has class conditional |
+
+ True
+ |
+
base distributions
+ +normflows/core.py
305 +306 +307 +308 +309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 |
|
forward(x, y=None)
+
+Get negative log-likelihood for maximum likelihood training
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch of data |
+ + required + | +
y |
+ + | +Batch of targets, if applicable |
+
+ None
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Negative log-likelihood of the batch |
+
normflows/core.py
337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 |
|
forward_kld(x, y=None)
+
+Estimates forward KL divergence, see see arXiv 1912.02762
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch sampled from target distribution |
+ + required + | +
y |
+ + | +Batch of targets, if applicable |
+
+ None
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Estimate of forward KL divergence averaged over batch |
+
normflows/core.py
325 +326 +327 +328 +329 +330 +331 +332 +333 +334 +335 |
|
load(path)
+
+Load model from state dict
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
path |
+ + | +Path including filename where to load model from |
+ + required + | +
normflows/core.py
422 +423 +424 +425 +426 +427 +428 |
|
log_prob(x, y)
+
+Get log probability for batch
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch |
+ + required + | +
y |
+ + | +Classes of x |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability |
+
normflows/core.py
384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 |
|
reset_temperature()
+
+Set temperature values of base distributions back to None
+ +normflows/core.py
445 +446 +447 +448 +449 |
|
sample(num_samples=1, y=None, temperature=None)
+
+Samples from flow-based approximate distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw |
+
+ 1
+ |
+
y |
+ + | +Classes to sample from, will be sampled uniformly if None |
+
+ None
+ |
+
temperature |
+ + | +Temperature parameter for temp annealed sampling |
+
+ None
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples, log probability |
+
normflows/core.py
349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 |
|
save(path)
+
+Save state dict of model
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
path |
+ + | +Path including filename where to save model |
+ + required + | +
normflows/core.py
414 +415 +416 +417 +418 +419 +420 |
|
set_temperature(temperature)
+
+Set temperature for temperature a annealed sampling
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
temperature |
+ + | +Temperature parameter |
+ + required + | +
normflows/core.py
430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 |
|
NormalizingFlow
+
+
+
+ Bases: nn.Module
Normalizing Flow model to approximate target distribution
+ + +normflows/core.py
9 + 10 + 11 + 12 + 13 + 14 + 15 + 16 + 17 + 18 + 19 + 20 + 21 + 22 + 23 + 24 + 25 + 26 + 27 + 28 + 29 + 30 + 31 + 32 + 33 + 34 + 35 + 36 + 37 + 38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 |
|
__init__(q0, flows, p=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
q0 |
+ + | +Base distribution |
+ + required + | +
flows |
+ + | +List of flows |
+ + required + | +
p |
+ + | +Target distribution |
+
+ None
+ |
+
normflows/core.py
14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 |
|
forward(z)
+
+Transforms latent variable z to the flow variable x
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Batch in the latent space |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Batch in the space of the target distribution |
+
normflows/core.py
27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 |
|
forward_and_log_det(z)
+
+Transforms latent variable z to the flow variable x and +computes log determinant of the Jacobian
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Batch in the latent space |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Batch in the space of the target distribution, |
+
+ | +log determinant of the Jacobian |
+
normflows/core.py
40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 |
|
forward_kld(x)
+
+Estimates forward KL divergence, see arXiv 1912.02762
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch sampled from target distribution |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Estimate of forward KL divergence averaged over batch |
+
normflows/core.py
87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 |
|
inverse(x)
+
+Transforms flow variable x to the latent variable z
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch in the space of the target distribution |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Batch in the latent space |
+
normflows/core.py
57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 |
|
inverse_and_log_det(x)
+
+Transforms flow variable x to the latent variable z and +computes log determinant of the Jacobian
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch in the space of the target distribution |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Batch in the latent space, log determinant of the |
+
+ | +Jacobian |
+
normflows/core.py
70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 +84 +85 |
|
load(path)
+
+Load model from state dict
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
path |
+ + | +Path including filename where to load model from |
+ + required + | +
normflows/core.py
207 +208 +209 +210 +211 +212 +213 |
|
log_prob(x)
+
+Get log probability for batch
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Batch |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability |
+
normflows/core.py
182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 |
|
reverse_alpha_div(num_samples=1, alpha=1, dreg=False)
+
+Alpha divergence when sampling from q
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw |
+
+ 1
+ |
+
dreg |
+ + | +Flag whether to use Double Reparametrized Gradient estimator, see arXiv 1810.04152 |
+
+ False
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Alpha divergence |
+
normflows/core.py
133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 |
|
reverse_kld(num_samples=1, beta=1.0, score_fn=True)
+
+Estimates reverse KL divergence, see arXiv 1912.02762
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw from base distribution |
+
+ 1
+ |
+
beta |
+ + | +Annealing parameter, see arXiv 1505.05770 |
+
+ 1.0
+ |
+
score_fn |
+ + | +Flag whether to include score function in gradient, see arXiv 1703.09194 |
+
+ True
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Estimate of the reverse KL divergence averaged over latent samples |
+
normflows/core.py
104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 |
|
sample(num_samples=1)
+
+Samples from flow-based approximate distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples, log probability |
+
normflows/core.py
167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 |
|
save(path)
+
+Save state dict of model
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
path |
+ + | +Path including filename where to save model |
+ + required + | +
normflows/core.py
199 +200 +201 +202 +203 +204 +205 |
|
NormalizingFlowVAE
+
+
+
+ Bases: nn.Module
VAE using normalizing flows to express approximate distribution
+ + +normflows/core.py
452 +453 +454 +455 +456 +457 +458 +459 +460 +461 +462 +463 +464 +465 +466 +467 +468 +469 +470 +471 +472 +473 +474 +475 +476 +477 +478 +479 +480 +481 +482 +483 +484 +485 +486 +487 +488 +489 +490 +491 +492 +493 +494 +495 +496 |
|
__init__(prior, q0=distributions.Dirac(), flows=None, decoder=None)
+
+Constructor of normalizing flow model
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
prior |
+ + | +Prior distribution of te VAE, i.e. Gaussian |
+ + required + | +
decoder |
+ + | +Optional decoder |
+
+ None
+ |
+
flows |
+ + | +Flows to transform output of base encoder |
+
+ None
+ |
+
q0 |
+ + | +Base Encoder |
+
+ distributions.Dirac()
+ |
+
normflows/core.py
457 +458 +459 +460 +461 +462 +463 +464 +465 +466 +467 +468 +469 +470 |
|
forward(x, num_samples=1)
+
+Takes data batch, samples num_samples for each data point from base distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +data batch |
+ + required + | +
num_samples |
+ + | +number of samples to draw for each data point |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +latent variables for each batch and sample, log_q, and log_p |
+
normflows/core.py
472 +473 +474 +475 +476 +477 +478 +479 +480 +481 +482 +483 +484 +485 +486 +487 +488 +489 +490 +491 +492 +493 +494 +495 +496 |
|
core_test
+
+
+distributions
+
+
+base
+
+
+AffineGaussian
+
+
+
+ Bases: BaseDistribution
Diagonal Gaussian an affine constant transformation applied to it, +can be class conditional or not
+ + +normflows/distributions/base.py
421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 +450 +451 +452 +453 +454 +455 +456 +457 +458 +459 +460 +461 +462 +463 +464 +465 +466 +467 +468 +469 +470 +471 +472 +473 +474 +475 +476 +477 +478 +479 +480 +481 +482 +483 +484 +485 +486 +487 +488 +489 +490 +491 +492 +493 +494 +495 +496 +497 +498 +499 +500 +501 +502 +503 +504 +505 +506 +507 +508 +509 +510 +511 +512 +513 +514 +515 +516 +517 |
|
__init__(shape, affine_shape, num_classes=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Shape of the variables |
+ + required + | +
affine_shape |
+ + | +Shape of the parameters in the affine transformation |
+ + required + | +
num_classes |
+ + | +Number of classes if the base is class conditional, None otherwise |
+
+ None
+ |
+
normflows/distributions/base.py
427 +428 +429 +430 +431 +432 +433 +434 +435 +436 +437 +438 +439 +440 +441 +442 +443 +444 +445 +446 +447 +448 +449 +450 +451 +452 +453 |
|
BaseDistribution
+
+
+
+ Bases: nn.Module
Base distribution of a flow-based model +Parameters do not depend of target variable (as is the case for a VAE encoder)
+ + +normflows/distributions/base.py
8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 |
|
forward(num_samples=1)
+
+Samples from base distribution and calculates log probability
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw from the distriubtion |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples drawn from the distribution, log probability |
+
normflows/distributions/base.py
17 +18 +19 +20 +21 +22 +23 +24 +25 +26 |
|
log_prob(z)
+
+Calculate log probability of batch of samples
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Batch of random variables to determine log probability for |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability for each batch element |
+
normflows/distributions/base.py
28 +29 +30 +31 +32 +33 +34 +35 +36 +37 |
|
sample(num_samples=1, **kwargs)
+
+Samples from base distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw from the distriubtion |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples drawn from the distribution |
+
normflows/distributions/base.py
39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 |
|
ClassCondDiagGaussian
+
+
+
+ Bases: BaseDistribution
Class conditional multivariate Gaussian distribution with diagonal covariance matrix
+ + +normflows/distributions/base.py
220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 |
|
__init__(shape, num_classes)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Tuple with shape of data, if int shape has one dimension |
+ + required + | +
num_classes |
+ + | +Number of classes to condition on |
+ + required + | +
normflows/distributions/base.py
225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 |
|
DiagGaussian
+
+
+
+ Bases: BaseDistribution
Multivariate Gaussian distribution with diagonal covariance matrix
+ + +normflows/distributions/base.py
52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 |
|
__init__(shape, trainable=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Tuple with shape of data, if int shape has one dimension |
+ + required + | +
normflows/distributions/base.py
57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 |
|
GaussianMixture
+
+
+
+ Bases: BaseDistribution
Mixture of Gaussians with diagonal covariance matrix
+ + +normflows/distributions/base.py
520 +521 +522 +523 +524 +525 +526 +527 +528 +529 +530 +531 +532 +533 +534 +535 +536 +537 +538 +539 +540 +541 +542 +543 +544 +545 +546 +547 +548 +549 +550 +551 +552 +553 +554 +555 +556 +557 +558 +559 +560 +561 +562 +563 +564 +565 +566 +567 +568 +569 +570 +571 +572 +573 +574 +575 +576 +577 +578 +579 +580 +581 +582 +583 +584 +585 +586 +587 +588 +589 +590 +591 +592 +593 +594 +595 +596 +597 +598 +599 +600 +601 +602 +603 +604 +605 +606 |
|
__init__(n_modes, dim, loc=None, scale=None, weights=None, trainable=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
n_modes |
+ + | +Number of modes of the mixture model |
+ + required + | +
dim |
+ + | +Number of dimensions of each Gaussian |
+ + required + | +
loc |
+ + | +List of mean values |
+
+ None
+ |
+
scale |
+ + | +List of diagonals of the covariance matrices |
+
+ None
+ |
+
weights |
+ + | +List of mode probabilities |
+
+ None
+ |
+
trainable |
+ + | +Flag, if true parameters will be optimized during training |
+
+ True
+ |
+
normflows/distributions/base.py
525 +526 +527 +528 +529 +530 +531 +532 +533 +534 +535 +536 +537 +538 +539 +540 +541 +542 +543 +544 +545 +546 +547 +548 +549 +550 +551 +552 +553 +554 +555 +556 +557 +558 +559 +560 +561 |
|
GaussianPCA
+
+
+
+ Bases: BaseDistribution
Gaussian distribution resulting from linearly mapping a normal distributed latent +variable describing the "content of the target"
+ + +normflows/distributions/base.py
609 +610 +611 +612 +613 +614 +615 +616 +617 +618 +619 +620 +621 +622 +623 +624 +625 +626 +627 +628 +629 +630 +631 +632 +633 +634 +635 +636 +637 +638 +639 +640 +641 +642 +643 +644 +645 +646 +647 +648 +649 +650 +651 +652 +653 +654 +655 +656 +657 +658 +659 +660 +661 +662 +663 +664 +665 +666 |
|
__init__(dim, latent_dim=None, sigma=0.1)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
dim |
+ + | +Number of dimensions of the flow variables |
+ + required + | +
latent_dim |
+ + | +Number of dimensions of the latent "content" variable; + if None it is set equal to dim |
+
+ None
+ |
+
sigma |
+ + | +Noise level |
+
+ 0.1
+ |
+
normflows/distributions/base.py
615 +616 +617 +618 +619 +620 +621 +622 +623 +624 +625 +626 +627 +628 +629 +630 +631 +632 +633 +634 |
|
GlowBase
+
+
+
+ Bases: BaseDistribution
Base distribution of the Glow model, i.e. Diagonal Gaussian with one mean and +log scale for each channel
+ + +normflows/distributions/base.py
294 +295 +296 +297 +298 +299 +300 +301 +302 +303 +304 +305 +306 +307 +308 +309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 +324 +325 +326 +327 +328 +329 +330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 +343 +344 +345 +346 +347 +348 +349 +350 +351 +352 +353 +354 +355 +356 +357 +358 +359 +360 +361 +362 +363 +364 +365 +366 +367 +368 +369 +370 +371 +372 +373 +374 +375 +376 +377 +378 +379 +380 +381 +382 +383 +384 +385 +386 +387 +388 +389 +390 +391 +392 +393 +394 +395 +396 +397 +398 +399 +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 |
|
__init__(shape, num_classes=None, logscale_factor=3.0)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Shape of the variables |
+ + required + | +
num_classes |
+ + | +Number of classes if the base is class conditional, None otherwise |
+
+ None
+ |
+
logscale_factor |
+ + | +Scaling factor for mean and log variance |
+
+ 3.0
+ |
+
normflows/distributions/base.py
300 +301 +302 +303 +304 +305 +306 +307 +308 +309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 +324 +325 +326 +327 +328 +329 +330 +331 +332 +333 +334 +335 +336 +337 +338 +339 +340 +341 +342 |
|
Uniform
+
+
+
+ Bases: BaseDistribution
Multivariate uniform distribution
+ + +normflows/distributions/base.py
105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 |
|
__init__(shape, low=-1.0, high=1.0)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Tuple with shape of data, if int shape has one dimension |
+ + required + | +
low |
+ + | +Lower bound of uniform distribution |
+
+ -1.0
+ |
+
high |
+ + | +Upper bound of uniform distribution |
+
+ 1.0
+ |
+
normflows/distributions/base.py
110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 |
|
UniformGaussian
+
+
+
+ Bases: BaseDistribution
Distribution of a 1D random variable with some entries having a uniform and +others a Gaussian distribution
+ + +normflows/distributions/base.py
145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 |
|
__init__(ndim, ind, scale=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
ndim |
+ + | +Int, number of dimensions |
+ + required + | +
ind |
+ + | +Iterable, indices of uniformly distributed entries |
+ + required + | +
scale |
+ + | +Iterable, standard deviation of Gaussian or width of uniform distribution |
+
+ None
+ |
+
normflows/distributions/base.py
151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 |
|
base_test
+
+
+decoder
+
+
+BaseDecoder
+
+
+
+ Bases: nn.Module
normflows/distributions/decoder.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 |
|
forward(z)
+
+Decodes z to x
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +x, std of x |
+
normflows/distributions/decoder.py
10 +11 +12 +13 +14 +15 +16 +17 +18 +19 |
|
log_prob(x, z)
+
+Log probability
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +observable |
+ + required + | +
z |
+ + | +latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log(p) of x given z |
+
normflows/distributions/decoder.py
21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 |
|
NNBernoulliDecoder
+
+
+
+ Bases: BaseDecoder
BaseDecoder representing a Bernoulli distribution with mean parametrized by a NN
+ + +normflows/distributions/decoder.py
73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 |
|
__init__(net)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
net |
+ + | +neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) |
+ + required + | +
normflows/distributions/decoder.py
78 +79 +80 +81 +82 +83 +84 +85 |
|
NNDiagGaussianDecoder
+
+
+
+ Bases: BaseDecoder
BaseDecoder representing a diagonal Gaussian distribution with mean and std parametrized by a NN
+ + +normflows/distributions/decoder.py
34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 |
|
__init__(net)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
net |
+ + | +neural network parametrizing mean and standard deviation of diagonal Gaussian |
+ + required + | +
normflows/distributions/decoder.py
39 +40 +41 +42 +43 +44 +45 +46 |
|
decoder_test
+
+
+distribution_test
+
+
+DistributionTest
+
+
+
+ Bases: unittest.TestCase
Generic test case for distribution modules
+ + +normflows/distributions/distribution_test.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 |
|
encoder
+
+
+BaseEncoder
+
+
+
+ Bases: nn.Module
Base distribution of a flow-based variational autoencoder +Parameters of the distribution depend of the target variable x
+ + +normflows/distributions/encoder.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 |
|
forward(x, num_samples=1)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Variable to condition on, first dimension is batch size |
+ + required + | +
num_samples |
+ + | +number of samples to draw per element of mini-batch |
+
+ 1
+ |
+
Returns + sample of z for x, log probability for sample
+ +normflows/distributions/encoder.py
15 +16 +17 +18 +19 +20 +21 +22 +23 +24 |
|
log_prob(z, x)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Primary random variable, first dimension is batch size |
+ + required + | +
x |
+ + | +Variable to condition on, first dimension is batch size |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of z given x |
+
normflows/distributions/encoder.py
26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 |
|
ConstDiagGaussian
+
+
+
+ Bases: BaseEncoder
normflows/distributions/encoder.py
74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 |
|
__init__(loc, scale)
+
+Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
loc |
+ + | +mean vector of the distribution |
+ + required + | +
scale |
+ + | +vector of the standard deviations on the diagonal of the covariance matrix |
+ + required + | +
normflows/distributions/encoder.py
75 +76 +77 +78 +79 +80 +81 +82 +83 +84 +85 +86 +87 +88 +89 |
|
forward(x=None, num_samples=1)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Variable to condition on, will only be used to determine the batch size |
+
+ None
+ |
+
num_samples |
+ + | +number of samples to draw per element of mini-batch |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +sample of z for x, log probability for sample |
+
normflows/distributions/encoder.py
91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 |
|
log_prob(z, x)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Primary random variable, first dimension is batch dimension |
+ + required + | +
x |
+ + | +Variable to condition on, first dimension is batch dimension |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of z given x |
+
normflows/distributions/encoder.py
111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 |
|
NNDiagGaussian
+
+
+
+ Bases: BaseEncoder
Diagonal Gaussian distribution with mean and variance determined by a neural network
+ + +normflows/distributions/encoder.py
130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 |
|
__init__(net)
+
+Construtor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
net |
+ + | +net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) |
+ + required + | +
normflows/distributions/encoder.py
135 +136 +137 +138 +139 +140 +141 +142 |
|
forward(x, num_samples=1)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
x |
+ + | +Variable to condition on |
+ + required + | +
num_samples |
+ + | +number of samples to draw per element of mini-batch |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +sample of z for x, log probability for sample |
+
normflows/distributions/encoder.py
144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 |
|
log_prob(z, x)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Primary random variable, first dimension is batch dimension |
+ + required + | +
x |
+ + | +Variable to condition on, first dimension is batch dimension |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of z given x |
+
normflows/distributions/encoder.py
167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 |
|
encoder_test
+
+
+linear_interpolation
+
+
+LinearInterpolation
+
+
+Linear interpolation of two distributions in the log space
+ + +normflows/distributions/linear_interpolation.py
1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 |
|
__init__(dist1, dist2, alpha)
+
+Constructor
+Interpolation parameter alpha:
+log_p = alpha * log_p_1 + (1 - alpha) * log_p_2
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
dist1 |
+ + | +First distribution |
+ + required + | +
dist2 |
+ + | +Second distribution |
+ + required + | +
alpha |
+ + | +Interpolation parameter |
+ + required + | +
normflows/distributions/linear_interpolation.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 |
|
mh_proposal
+
+
+DiagGaussianProposal
+
+
+
+ Bases: MHProposal
Diagonal Gaussian distribution with previous value as mean +as a proposal for Metropolis Hastings algorithm
+ + +normflows/distributions/mh_proposal.py
47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 |
|
__init__(shape, scale)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Shape of variables to sample |
+ + required + | +
scale |
+ + | +Standard deviation of distribution |
+ + required + | +
normflows/distributions/mh_proposal.py
53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 |
|
MHProposal
+
+
+
+ Bases: nn.Module
Proposal distribution for the Metropolis Hastings algorithm
+ + +normflows/distributions/mh_proposal.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 |
|
forward(z)
+
+Draw samples given z and compute log probability difference
+log(p(z | z_new)) - log(p(z_new | z))
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +Previous samples |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Proposal, difference of log probability ratio |
+
normflows/distributions/mh_proposal.py
31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 |
|
log_prob(z_, z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z_ |
+ + | +Potential new sample |
+ + required + | +
z |
+ + | +Previous sample |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Log probability of proposal distribution |
+
normflows/distributions/mh_proposal.py
20 +21 +22 +23 +24 +25 +26 +27 +28 +29 |
|
sample(z)
+
+Sample new value based on previous z
+ +normflows/distributions/mh_proposal.py
14 +15 +16 +17 +18 |
|
prior
+
+
+ImagePrior
+
+
+
+ Bases: nn.Module
Intensities of an image determine probability density of prior
+ + +normflows/distributions/prior.py
21 + 22 + 23 + 24 + 25 + 26 + 27 + 28 + 29 + 30 + 31 + 32 + 33 + 34 + 35 + 36 + 37 + 38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 |
|
__init__(image, x_range=[-3, 3], y_range=[-3, 3], eps=1e-10)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
image |
+ + | +image as np matrix |
+ + required + | +
x_range |
+ + | +x range to position image at |
+
+ [-3, 3]
+ |
+
y_range |
+ + | +y range to position image at |
+
+ [-3, 3]
+ |
+
eps |
+ + | +small value to add to image to avoid log(0) problems |
+
+ 1e-10
+ |
+
normflows/distributions/prior.py
26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 |
|
log_prob(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 |
|
rejection_sampling(num_steps=1)
+
+Perform rejection sampling on image distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_steps |
+ + | +Number of rejection sampling steps to perform |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Accepted samples |
+
normflows/distributions/prior.py
71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 +84 +85 +86 +87 +88 |
|
sample(num_samples=1)
+
+Sample from image distribution through rejection sampling
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples |
+
normflows/distributions/prior.py
90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 |
|
PriorDistribution
+
+
+normflows/distributions/prior.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 |
|
log_prob(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
10 +11 +12 +13 +14 +15 +16 +17 +18 |
|
Sinusoidal
+
+
+
+ Bases: PriorDistribution
normflows/distributions/prior.py
152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 |
|
__init__(scale, period)
+
+Distribution 2d with sinusoidal density +given by
+w_1(z) = sin(2*pi / period * z[0])
+log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
scale |
+ + | +scale of the distribution, see formula |
+ + required + | +
period |
+ + | +period of the sinosoidal |
+ + required + | +
normflows/distributions/prior.py
153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 |
|
log_prob(z)
+
+log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2
+w_1(z) = sin(2*pi / period * z[0])
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 |
|
Sinusoidal_gap
+
+
+
+ Bases: PriorDistribution
normflows/distributions/prior.py
197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 |
|
__init__(scale, period)
+
+Distribution 2d with sinusoidal density with gap +given by
+w_1(z) = sin(2*pi / period * z[0])
+w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2)
+log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2))
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
loc |
+ + | +distance of modes from the origin |
+ + required + | +
scale |
+ + | +scale of modes |
+ + required + | +
normflows/distributions/prior.py
198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 |
|
log_prob(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 |
|
Sinusoidal_split
+
+
+
+ Bases: PriorDistribution
normflows/distributions/prior.py
248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 +292 +293 +294 +295 +296 |
|
__init__(scale, period)
+
+Distribution 2d with sinusoidal density with split +given by
+w_1(z) = sin(2*pi / period * z[0])
+w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3)
+log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2))
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
loc |
+ + | +distance of modes from the origin |
+ + required + | +
scale |
+ + | +scale of modes |
+ + required + | +
normflows/distributions/prior.py
249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 |
|
log_prob(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 +292 +293 +294 +295 +296 |
|
Smiley
+
+
+
+ Bases: PriorDistribution
normflows/distributions/prior.py
299 +300 +301 +302 +303 +304 +305 +306 +307 +308 +309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 +324 +325 +326 +327 |
|
__init__(scale)
+
+Distribution 2d of a smiley :)
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
scale |
+ + | +scale of the smiley |
+ + required + | +
normflows/distributions/prior.py
300 +301 +302 +303 +304 +305 +306 +307 |
|
log_prob(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
309 +310 +311 +312 +313 +314 +315 +316 +317 +318 +319 +320 +321 +322 +323 +324 +325 +326 +327 |
|
TwoModes
+
+
+
+ Bases: PriorDistribution
normflows/distributions/prior.py
107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 |
|
__init__(loc, scale)
+
+Distribution 2d with two modes
+Distribution 2d with two modes at
+z[0] = -loc
and z[0] = loc
+following the density
log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2
+ - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2))
+
+Args: + loc: distance of modes from the origin + scale: scale of modes
+ +normflows/distributions/prior.py
108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 |
|
log_prob(z)
+
+log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2
+ - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2))
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/prior.py
126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 |
|
prior_test
+
+
+target
+
+
+CircularGaussianMixture
+
+
+
+ Bases: nn.Module
Two-dimensional Gaussian mixture arranged in a circle
+ + +normflows/distributions/target.py
132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 |
|
__init__(n_modes=8)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
n_modes |
+ + | +Number of modes |
+
+ 8
+ |
+
normflows/distributions/target.py
137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 |
|
RingMixture
+
+
+
+ Bases: Target
Mixture of ring distributions in two dimensions
+ + +normflows/distributions/target.py
176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 |
|
Target
+
+
+
+ Bases: nn.Module
Sample target distributions to test models
+ + +normflows/distributions/target.py
8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 |
|
__init__(prop_scale=torch.tensor(6.0), prop_shift=torch.tensor(-3.0))
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
prop_scale |
+ + | +Scale for the uniform proposal |
+
+ torch.tensor(6.0)
+ |
+
prop_shift |
+ + | +Shift for the uniform proposal |
+
+ torch.tensor(-3.0)
+ |
+
normflows/distributions/target.py
13 +14 +15 +16 +17 +18 +19 +20 +21 +22 |
|
log_prob(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/target.py
24 +25 +26 +27 +28 +29 +30 +31 +32 |
|
rejection_sampling(num_steps=1)
+
+Perform rejection sampling on image distribution
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_steps |
+ + | +Number of rejection sampling steps to perform |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Accepted samples |
+
normflows/distributions/target.py
34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 |
|
sample(num_samples=1)
+
+Sample from image distribution through rejection sampling
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +Number of samples to draw |
+
+ 1
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Samples |
+
normflows/distributions/target.py
57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 |
|
TwoIndependent
+
+
+
+ Bases: Target
Target distribution that combines two independent distributions of equal +size into one distribution. This is needed for Augmented Normalizing Flows, +see https://arxiv.org/abs/2002.07101
+ + +normflows/distributions/target.py
76 +77 +78 +79 +80 +81 +82 +83 +84 +85 +86 +87 +88 +89 +90 +91 +92 +93 +94 +95 +96 |
|
TwoMoons
+
+
+
+ Bases: Target
Bimodal two-dimensional distribution
+ + +normflows/distributions/target.py
99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 |
|
log_prob(z)
+
+log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2
+ + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2)
+ + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2))
+
+
+ Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +value or batch of latent variable |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +log probability of the distribution for z |
+
normflows/distributions/target.py
109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 |
|
target_test
+
+
+flows
+
+
+affine
+
+
+autoregressive
+
+
+Autoregressive
+
+
+
+ Bases: Flow
Transforms each input variable with an invertible elementwise transformation.
+The parameters of each invertible elementwise transformation can be functions of previous input +variables, but they must not depend on the current or any following input variables.
+NOTE Calculating the inverse transform is D times slower than calculating the +forward transform, where D is the dimensionality of the input to the transform.
+ + +normflows/flows/affine/autoregressive.py
10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 |
|
autoregressive_test
+
+
+coupling
+
+
+AffineConstFlow
+
+
+
+ Bases: Flow
scales and shifts with learned constants per dimension. In the NICE paper there is a +scaling layer which is a special case of this where t is None
+ + +normflows/flows/affine/coupling.py
9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 |
|
__init__(shape, scale=True, shift=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Shape of the coupling layer |
+ + required + | +
scale |
+ + | +Flag whether to apply scaling |
+
+ True
+ |
+
shift |
+ + | +Flag whether to apply shift |
+
+ True
+ |
+
logscale_factor |
+ + | +Optional factor which can be used to control the scale of the log scale factor |
+ + required + | +
normflows/flows/affine/coupling.py
15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 |
|
AffineCoupling
+
+
+
+ Bases: Flow
Affine Coupling layer as introduced RealNVP paper, see arXiv: 1605.08803
+ + +normflows/flows/affine/coupling.py
99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 |
|
__init__(param_map, scale=True, scale_map='exp')
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
param_map |
+ + | +Maps features to shift and scale parameter (if applicable) |
+ + required + | +
scale |
+ + | +Flag whether scale shall be applied |
+
+ True
+ |
+
scale_map |
+ + | +Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model |
+
+ 'exp'
+ |
+
normflows/flows/affine/coupling.py
104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 |
|
forward(z)
+
+z is a list of z1 and z2; z = [z1, z2]
+z1 is left constant and affine map is applied to z2 with parameters depending
+on z1
normflows/flows/affine/coupling.py
117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 |
|
AffineCouplingBlock
+
+
+
+ Bases: Flow
Affine Coupling layer including split and merge operation
+ + +normflows/flows/affine/coupling.py
232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 |
|
__init__(param_map, scale=True, scale_map='exp', split_mode='channel')
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
param_map |
+ + | +Maps features to shift and scale parameter (if applicable) |
+ + required + | +
scale |
+ + | +Flag whether scale shall be applied |
+
+ True
+ |
+
scale_map |
+ + | +Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow |
+
+ 'exp'
+ |
+
split_mode |
+ + | +Splitting mode, for possible values see Split class |
+
+ 'channel'
+ |
+
normflows/flows/affine/coupling.py
237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 |
|
CCAffineConst
+
+
+
+ Bases: Flow
Affine constant flow layer with class-conditional parameters
+ + +normflows/flows/affine/coupling.py
57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 +84 +85 +86 +87 +88 +89 +90 +91 +92 +93 +94 +95 +96 |
|
MaskedAffineFlow
+
+
+
+ Bases: Flow
RealNVP as introduced in arXiv: 1605.08803
+Masked affine flow:
+f(z) = b * z + (1 - b) * (z * exp(s(b * z)) + t)
+
+normflows/flows/affine/coupling.py
174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 |
|
__init__(b, t=None, s=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
b |
+ + | +mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s |
+ + required + | +
t |
+ + | +translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied |
+
+ None
+ |
+
s |
+ + | +scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied |
+
+ None
+ |
+
normflows/flows/affine/coupling.py
187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 |
|
coupling_test
+
+
+glow
+
+
+GlowBlock
+
+
+
+ Bases: Flow
Glow: Generative Flow with Invertible 1×1 Convolutions, arXiv: 1807.03039
+One Block of the Glow model, comprised of
+normflows/flows/affine/glow.py
11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 +84 |
|
__init__(channels, hidden_channels, scale=True, scale_map='sigmoid', split_mode='channel', leaky=0.0, init_zeros=True, use_lu=True, net_actnorm=False)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
channels |
+ + | +Number of channels of the data |
+ + required + | +
hidden_channels |
+ + | +number of channels in the hidden layer of the ConvNet |
+ + required + | +
scale |
+ + | +Flag, whether to include scale in affine coupling layer |
+
+ True
+ |
+
scale_map |
+ + | +Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow |
+
+ 'sigmoid'
+ |
+
split_mode |
+ + | +Splitting mode, for possible values see Split class |
+
+ 'channel'
+ |
+
leaky |
+ + | +Leaky parameter of LeakyReLUs of ConvNet2d |
+
+ 0.0
+ |
+
init_zeros |
+ + | +Flag whether to initialize last conv layer with zeros |
+
+ True
+ |
+
use_lu |
+ + | +Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers |
+
+ True
+ |
+
logscale_factor |
+ + | +Factor which can be used to control the scale of the log scale factor, see source |
+ + required + | +
normflows/flows/affine/glow.py
21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 |
|
glow_test
+
+
+base
+
+
+Composite
+
+
+
+ Bases: Flow
Composes several flows into one, in the order they are given.
+ + +normflows/flows/base.py
48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 |
|
__init__(flows)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
flows |
+ + | +Iterable of flows to composite |
+ + required + | +
normflows/flows/base.py
53 +54 +55 +56 +57 +58 +59 +60 |
|
Flow
+
+
+
+ Bases: nn.Module
Generic class for flow functions
+ + +normflows/flows/base.py
5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 |
|
forward(z)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
z |
+ + | +input variable, first dimension is batch dim |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +transformed z and log of absolute determinant |
+
normflows/flows/base.py
13 +14 +15 +16 +17 +18 +19 +20 +21 |
|
Reverse
+
+
+
+ Bases: Flow
Switches the forward transform of a flow layer with its inverse and vice versa
+ + +normflows/flows/base.py
27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 |
|
__init__(flow)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
flow |
+ + | +Flow layer to be reversed |
+ + required + | +
normflows/flows/base.py
32 +33 +34 +35 +36 +37 +38 +39 |
|
base_test
+
+
+flow_test
+
+
+FlowTest
+
+
+
+ Bases: unittest.TestCase
Generic test case for flow modules
+ + +normflows/flows/flow_test.py
7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 |
|
mixing
+
+
+Invertible1x1Conv
+
+
+
+ Bases: Flow
Invertible 1x1 convolution introduced in the Glow paper +Assumes 4d input/output tensors of the form NCHW
+ + +normflows/flows/mixing.py
57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 |
|
__init__(num_channels, use_lu=False)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_channels |
+ + | +Number of channels of the data |
+ + required + | +
use_lu |
+ + | +Flag whether to parametrize weights through the LU decomposition |
+
+ False
+ |
+
normflows/flows/mixing.py
63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 +84 +85 +86 |
|
InvertibleAffine
+
+
+
+ Bases: Flow
Invertible affine transformation without shift, i.e. one-dimensional +version of the invertible 1x1 convolutions
+ + +normflows/flows/mixing.py
136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 |
|
__init__(num_channels, use_lu=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_channels |
+ + | +Number of channels of the data |
+ + required + | +
use_lu |
+ + | +Flag whether to parametrize weights through the LU decomposition |
+
+ True
+ |
+
normflows/flows/mixing.py
142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 |
|
LULinearPermute
+
+
+
+ Bases: Flow
Fixed permutation combined with a linear transformation parametrized +using the LU decomposition, used in https://arxiv.org/abs/1906.04032
+ + +normflows/flows/mixing.py
535 +536 +537 +538 +539 +540 +541 +542 +543 +544 +545 +546 +547 +548 +549 +550 +551 +552 +553 +554 +555 +556 +557 +558 +559 +560 +561 +562 +563 |
|
__init__(num_channels, identity_init=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_channels |
+ + | +Number of dimensions of the data |
+ + required + | +
identity_init |
+ + | +Flag, whether to initialize linear transform as identity matrix |
+
+ True
+ |
+
normflows/flows/mixing.py
541 +542 +543 +544 +545 +546 +547 +548 +549 +550 +551 +552 +553 |
|
Permute
+
+
+
+ Bases: Flow
Permutation features along the channel dimension
+ + +normflows/flows/mixing.py
9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 |
|
__init__(num_channels, mode='shuffle')
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_channel |
+ + | +Number of channels |
+ + required + | +
mode |
+ + | +Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part |
+
+ 'shuffle'
+ |
+
normflows/flows/mixing.py
14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 |
|
mixing_test
+
+
+neural_spline
+
+
+autoregressive
+
+
+Implementations of autoregressive transforms. +Code taken from https://github.com/bayesiains/nsf
+ + + +autoregressive_test
+
+
+Tests for the autoregressive transforms. +Code partially taken from https://github.com/bayesiains/nsf
+ + + +coupling
+
+
+Implementations of various coupling layers. +Code taken from https://github.com/bayesiains/nsf
+ + + +Coupling
+
+
+
+ Bases: Flow
A base class for coupling layers. Supports 2D inputs (NxD), as well as 4D inputs for +images (NxCxHxW). For images the splitting is done on the channel dimension, using the +provided 1D mask.
+ + +normflows/flows/neural_spline/coupling.py
16 + 17 + 18 + 19 + 20 + 21 + 22 + 23 + 24 + 25 + 26 + 27 + 28 + 29 + 30 + 31 + 32 + 33 + 34 + 35 + 36 + 37 + 38 + 39 + 40 + 41 + 42 + 43 + 44 + 45 + 46 + 47 + 48 + 49 + 50 + 51 + 52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 |
|
__init__(mask, transform_net_create_fn, unconditional_transform=None)
+
+Constructor.
+mask: a 1-dim tensor, tuple or list. It indexes inputs as follows:
+mask[i] > 0
, input[i]
will be transformed.mask[i] <= 0
, input[i]
will be passed unchanged.normflows/flows/neural_spline/coupling.py
21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 |
|
coupling_test
+
+
+Tests for the coupling Transforms. +Code partially taken from https://github.com/bayesiains/nsf
+ + + +wrapper
+
+
+AutoregressiveRationalQuadraticSpline
+
+
+
+ Bases: Flow
Neural spline flow coupling layer, wrapper for the implementation +of Durkan et al., see sources
+ + +normflows/flows/neural_spline/wrapper.py
174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 |
|
__init__(num_input_channels, num_blocks, num_hidden_channels, num_bins=8, tail_bound=3, activation=nn.ReLU, dropout_probability=0.0, permute_mask=False, init_identity=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_input_channels |
+
+ int
+ |
+ Flow dimension |
+ + required + | +
num_blocks |
+
+ int
+ |
+ Number of residual blocks of the parameter NN |
+ + required + | +
num_hidden_channels |
+
+ int
+ |
+ Number of hidden units of the NN |
+ + required + | +
num_bins |
+
+ int
+ |
+ Number of bins |
+
+ 8
+ |
+
tail_bound |
+
+ int
+ |
+ Bound of the spline tails |
+
+ 3
+ |
+
activation |
+
+ torch module
+ |
+ Activation function |
+
+ nn.ReLU
+ |
+
dropout_probability |
+
+ float
+ |
+ Dropout probability of the NN |
+
+ 0.0
+ |
+
permute_mask |
+
+ bool
+ |
+ Flag, permutes the mask of the NN |
+
+ False
+ |
+
init_identity |
+
+ bool
+ |
+ Flag, initialize transform as identity |
+
+ True
+ |
+
normflows/flows/neural_spline/wrapper.py
180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 |
|
CircularAutoregressiveRationalQuadraticSpline
+
+
+
+ Bases: Flow
Neural spline flow coupling layer, wrapper for the implementation +of Durkan et al., see sources
+ + +normflows/flows/neural_spline/wrapper.py
233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 +292 +293 +294 +295 |
|
__init__(num_input_channels, num_blocks, num_hidden_channels, ind_circ, num_bins=8, tail_bound=3, activation=nn.ReLU, dropout_probability=0.0, permute_mask=True, init_identity=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_input_channels |
+
+ int
+ |
+ Flow dimension |
+ + required + | +
num_blocks |
+
+ int
+ |
+ Number of residual blocks of the parameter NN |
+ + required + | +
num_hidden_channels |
+
+ int
+ |
+ Number of hidden units of the NN |
+ + required + | +
ind_circ |
+
+ Iterable
+ |
+ Indices of the circular coordinates |
+ + required + | +
num_bins |
+
+ int
+ |
+ Number of bins |
+
+ 8
+ |
+
tail_bound |
+
+ int
+ |
+ Bound of the spline tails |
+
+ 3
+ |
+
activation |
+
+ torch module
+ |
+ Activation function |
+
+ nn.ReLU
+ |
+
dropout_probability |
+
+ float
+ |
+ Dropout probability of the NN |
+
+ 0.0
+ |
+
permute_mask |
+
+ bool
+ |
+ Flag, permutes the mask of the NN |
+
+ True
+ |
+
init_identity |
+
+ bool
+ |
+ Flag, initialize transform as identity |
+
+ True
+ |
+
normflows/flows/neural_spline/wrapper.py
239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 |
|
CircularCoupledRationalQuadraticSpline
+
+
+
+ Bases: Flow
Neural spline flow coupling layer with circular coordinates
+ + +normflows/flows/neural_spline/wrapper.py
78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 |
|
__init__(num_input_channels, num_blocks, num_hidden_channels, ind_circ, num_bins=8, tail_bound=3.0, activation=nn.ReLU, dropout_probability=0.0, reverse_mask=False, mask=None, init_identity=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_input_channels |
+
+ int
+ |
+ Flow dimension |
+ + required + | +
num_blocks |
+
+ int
+ |
+ Number of residual blocks of the parameter NN |
+ + required + | +
num_hidden_channels |
+
+ int
+ |
+ Number of hidden units of the NN |
+ + required + | +
ind_circ |
+
+ Iterable
+ |
+ Indices of the circular coordinates |
+ + required + | +
num_bins |
+
+ int
+ |
+ Number of bins |
+
+ 8
+ |
+
tail_bound |
+
+ float or Iterable
+ |
+ Bound of the spline tails |
+
+ 3.0
+ |
+
activation |
+
+ torch module
+ |
+ Activation function |
+
+ nn.ReLU
+ |
+
dropout_probability |
+
+ float
+ |
+ Dropout probability of the NN |
+
+ 0.0
+ |
+
reverse_mask |
+
+ bool
+ |
+ Flag whether the reverse mask should be used |
+
+ False
+ |
+
mask |
+
+ torch tensor
+ |
+ Mask to be used, alternating masked generated is None |
+
+ None
+ |
+
init_identity |
+
+ bool
+ |
+ Flag, initialize transform as identity |
+
+ True
+ |
+
normflows/flows/neural_spline/wrapper.py
83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 |
|
CoupledRationalQuadraticSpline
+
+
+
+ Bases: Flow
Neural spline flow coupling layer, wrapper for the implementation +of Durkan et al., see source
+ + +normflows/flows/neural_spline/wrapper.py
14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 |
|
__init__(num_input_channels, num_blocks, num_hidden_channels, num_bins=8, tails='linear', tail_bound=3.0, activation=nn.ReLU, dropout_probability=0.0, reverse_mask=False)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_input_channels |
+
+ int
+ |
+ Flow dimension |
+ + required + | +
num_blocks |
+
+ int
+ |
+ Number of residual blocks of the parameter NN |
+ + required + | +
num_hidden_channels |
+
+ int
+ |
+ Number of hidden units of the NN |
+ + required + | +
num_bins |
+
+ int
+ |
+ Number of bins |
+
+ 8
+ |
+
tails |
+
+ str
+ |
+ Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval |
+
+ 'linear'
+ |
+
tail_bound |
+
+ float
+ |
+ Bound of the spline tails |
+
+ 3.0
+ |
+
activation |
+
+ torch module
+ |
+ Activation function |
+
+ nn.ReLU
+ |
+
dropout_probability |
+
+ float
+ |
+ Dropout probability of the NN |
+
+ 0.0
+ |
+
reverse_mask |
+
+ bool
+ |
+ Flag whether the reverse mask should be used |
+
+ False
+ |
+
normflows/flows/neural_spline/wrapper.py
20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 |
|
wrapper_test
+
+
+normalization
+
+
+ActNorm
+
+
+
+ Bases: AffineConstFlow
An AffineConstFlow but with a data-dependent initialization, +where on the very first batch we clever initialize the s,t so that the output +is unit gaussian. As described in Glow paper.
+ + +normflows/flows/normalization.py
7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 |
|
BatchNorm
+
+
+
+ Bases: Flow
Batch Normalization with out considering the derivatives of the batch statistics, see arXiv: 1605.08803
+ + +normflows/flows/normalization.py
42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 |
|
forward(z)
+
+Do batch norm over batch and sample dimension
+ +normflows/flows/normalization.py
52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 |
|
periodic
+
+
+PeriodicShift
+
+
+
+ Bases: Flow
Shift and wrap periodic coordinates
+ + +normflows/flows/periodic.py
35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 |
|
__init__(ind, bound=1.0, shift=0.0)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
ind |
+ + | +Iterable, indices of coordinates to be mapped |
+ + required + | +
bound |
+ + | +Float or iterable, bound of interval |
+
+ 1.0
+ |
+
shift |
+ + | +Tensor, shift to be applied |
+
+ 0.0
+ |
+
normflows/flows/periodic.py
40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 |
|
PeriodicWrap
+
+
+
+ Bases: Flow
Map periodic coordinates to fixed interval
+ + +normflows/flows/periodic.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 |
|
__init__(ind, bound=1.0)
+
+Constructor
+ind: Iterable, indices of coordinates to be mapped +bound: Float or iterable, bound of interval
+ +normflows/flows/periodic.py
11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 |
|
periodic_test
+
+
+planar
+
+
+Planar
+
+
+
+ Bases: Flow
Planar flow as introduced in arXiv: 1505.05770
+ f(z) = z + u * h(w * z + b)
+
+
+
+ normflows/flows/planar.py
8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 |
|
__init__(shape, act='tanh', u=None, w=None, b=None)
+
+Constructor of the planar flow
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +shape of the latent variable z |
+ + required + | +
h |
+ + | +nonlinear function h of the planar flow (see definition of f above) |
+ + required + | +
u,w,b |
+ + | +optional initialization for parameters |
+ + required + | +
normflows/flows/planar.py
16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 |
|
planar_test
+
+
+radial
+
+
+Radial
+
+
+
+ Bases: Flow
Radial flow as introduced in arXiv: 1505.05770
+ f(z) = z + beta * h(alpha, r) * (z - z_0)
+
+
+
+ normflows/flows/radial.py
8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 |
|
__init__(shape, z_0=None)
+
+Constructor of the radial flow
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +shape of the latent variable z |
+ + required + | +
z_0 |
+ + | +parameter of the radial flow |
+
+ None
+ |
+
normflows/flows/radial.py
16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 |
|
radial_test
+
+
+reshape
+
+
+Merge
+
+
+
+ Bases: Split
Same as Split but with forward and backward pass interchanged
+ + +normflows/flows/reshape.py
88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 |
|
Split
+
+
+
+ Bases: Flow
Split features into two sets
+ + +normflows/flows/reshape.py
9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 +82 +83 +84 +85 |
|
__init__(mode='channel')
+
+Constructor
+The splitting mode can be:
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
mode |
+ + | +splitting mode |
+
+ 'channel'
+ |
+
normflows/flows/reshape.py
14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 |
|
Squeeze
+
+
+
+ Bases: Flow
Squeeze operation of multi-scale architecture, RealNVP or Glow paper
+ + +normflows/flows/reshape.py
103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 |
|
__init__()
+
+Constructor
+ +normflows/flows/reshape.py
108 +109 +110 +111 +112 |
|
residual
+
+
+Residual
+
+
+
+ Bases: Flow
Invertible residual net block, wrapper to the implementation of Chen et al., +see sources
+ + +normflows/flows/residual.py
12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 |
|
__init__(net, reverse=True, reduce_memory=True, geom_p=0.5, lamb=2.0, n_power_series=None, exact_trace=False, brute_force=False, n_samples=1, n_exact_terms=2, n_dist='geometric')
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
net |
+ + | +Neural network, must be Lipschitz continuous with L < 1 |
+ + required + | +
reverse |
+ + | +Flag, if true the map |
+
+ True
+ |
+
reduce_memory |
+ + | +Flag, if true Neumann series and precomputations, for backward pass in forward pass are done |
+
+ True
+ |
+
geom_p |
+ + | +Parameter of the geometric distribution used for the Neumann series |
+
+ 0.5
+ |
+
lamb |
+ + | +Parameter of the geometric distribution used for the Neumann series |
+
+ 2.0
+ |
+
n_power_series |
+ + | +Number of terms in the Neumann series |
+
+ None
+ |
+
exact_trace |
+ + | +Flag, if true the trace of the Jacobian is computed exactly |
+
+ False
+ |
+
brute_force |
+ + | +Flag, if true the Jacobian is computed exactly in 2D |
+
+ False
+ |
+
n_samples |
+ + | +Number of samples used to estimate power series |
+
+ 1
+ |
+
n_exact_terms |
+ + | +Number of terms always included in the power series |
+
+ 2
+ |
+
n_dist |
+ + | +Distribution used for the power series, either "geometric" or "poisson" |
+
+ 'geometric'
+ |
+
normflows/flows/residual.py
18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 |
|
iResBlock
+
+
+
+ Bases: nn.Module
normflows/flows/residual.py
78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 +215 +216 +217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 |
|
__init__(nnet, geom_p=0.5, lamb=2.0, n_power_series=None, exact_trace=False, brute_force=False, n_samples=1, n_exact_terms=2, n_dist='geometric', neumann_grad=True, grad_in_forward=False)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
nnet |
+ + | +a nn.Module |
+ + required + | +
n_power_series |
+ + | +number of power series. If not None, uses a biased approximation to logdet. |
+
+ None
+ |
+
exact_trace |
+ + | +if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. |
+
+ False
+ |
+
brute_force |
+ + | +Computes the exact logdet. Only available for 2D inputs. |
+
+ False
+ |
+
normflows/flows/residual.py
79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 |
|
residual_test
+
+
+stochastic
+
+
+HamiltonianMonteCarlo
+
+
+
+ Bases: Flow
Flow layer using the HMC proposal in Stochastic Normalising Flows
+ + + +normflows/flows/stochastic.py
52 + 53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 |
|
__init__(target, steps, log_step_size, log_mass, max_abs_grad=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
target |
+ + | +The stationary distribution of this Markov transition, i.e. the target distribution to sample from. |
+ + required + | +
steps |
+ + | +The number of leapfrog steps |
+ + required + | +
log_step_size |
+ + | +The log step size used in the leapfrog integrator. shape (dim) |
+ + required + | +
log_mass |
+ + | +The log_mass determining the variance of the momentum samples. shape (dim) |
+ + required + | +
max_abs_grad |
+ + | +Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability. |
+
+ None
+ |
+
normflows/flows/stochastic.py
58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 |
|
MetropolisHastings
+
+
+
+ Bases: Flow
Sampling through Metropolis Hastings in Stochastic Normalizing Flow
+ + + +normflows/flows/stochastic.py
6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 |
|
__init__(target, proposal, steps)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
target |
+ + | +The stationary distribution of this Markov transition, i.e. the target distribution to sample from. |
+ + required + | +
proposal |
+ + | +Proposal distribution |
+ + required + | +
steps |
+ + | +Number of MCMC steps to perform |
+ + required + | +
normflows/flows/stochastic.py
12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 |
|
stochastic_test
+
+
+nets
+
+
+cnn
+
+
+ConvNet2d
+
+
+
+ Bases: nn.Module
Convolutional Neural Network with leaky ReLU nonlinearities
+ + +normflows/nets/cnn.py
5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 |
|
__init__(channels, kernel_size, leaky=0.0, init_zeros=True, actnorm=False, weight_std=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
channels |
+ + | +List of channels of conv layers, first entry is in_channels |
+ + required + | +
kernel_size |
+ + | +List of kernel sizes, same for height and width |
+ + required + | +
leaky |
+ + | +Leaky part of ReLU |
+
+ 0.0
+ |
+
init_zeros |
+ + | +Flag whether last layer shall be initialized with zeros |
+
+ True
+ |
+
scale_output |
+ + | +Flag whether to scale output with a log scale parameter |
+ + required + | +
logscale_factor |
+ + | +Constant factor to be multiplied to log scaling |
+ + required + | +
actnorm |
+ + | +Flag whether activation normalization shall be done after each conv layer except output |
+
+ False
+ |
+
weight_std |
+ + | +Fixed std used to initialize every layer |
+
+ None
+ |
+
normflows/nets/cnn.py
10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 |
|
lipschitz
+
+
+LipschitzCNN
+
+
+
+ Bases: nn.Module
Convolutional neural network which is Lipschitz continuous +with Lipschitz constant L < 1
+ + +normflows/nets/lipschitz.py
70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 |
|
__init__(channels, kernel_size, lipschitz_const=0.97, max_lipschitz_iter=5, lipschitz_tolerance=None, init_zeros=True)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
channels |
+ + | +Integer list with the number of channels of the layers |
+ + required + | +
kernel_size |
+ + | +Integer list of kernel sizes of the layers |
+ + required + | +
lipschitz_const |
+ + | +Maximum Lipschitz constant of each layer |
+
+ 0.97
+ |
+
max_lipschitz_iter |
+ + | +Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used |
+
+ 5
+ |
+
lipschitz_tolerance |
+ + | +Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 |
+
+ None
+ |
+
init_zeros |
+ + | +Flag, whether to initialize last layer approximately with zeros |
+
+ True
+ |
+
normflows/nets/lipschitz.py
76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 |
|
LipschitzMLP
+
+
+
+ Bases: nn.Module
Fully connected neural net which is Lipschitz continuou with Lipschitz constant L < 1
+ + +normflows/nets/lipschitz.py
14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 |
|
__init__(channels, lipschitz_const=0.97, max_lipschitz_iter=5, lipschitz_tolerance=None, init_zeros=True)
+
+Constructor + channels: Integer list with the number of channels of +the layers + lipschitz_const: Maximum Lipschitz constant of each layer + max_lipschitz_iter: Maximum number of iterations used to +ensure that layers are Lipschitz continuous with L smaller than +set maximum; if None, tolerance is used + lipschitz_tolerance: Float, tolerance used to ensure +Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 + init_zeros: Flag, whether to initialize last layer +approximately with zeros
+ +normflows/nets/lipschitz.py
17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 |
|
projmax_(v)
+
+Inplace argmax on absolute value.
+ +normflows/nets/lipschitz.py
651 +652 +653 +654 +655 +656 |
|
made
+
+
+Implementation of MADE. +Code taken from https://github.com/bayesiains/nsf
+ + + +MADE
+
+
+
+ Bases: nn.Module
Implementation of MADE.
+It can use either feedforward blocks or residual blocks (default is residual). +Optionally, it can use batch norm or dropout within blocks (default is no).
+ + +normflows/nets/made.py
217 +218 +219 +220 +221 +222 +223 +224 +225 +226 +227 +228 +229 +230 +231 +232 +233 +234 +235 +236 +237 +238 +239 +240 +241 +242 +243 +244 +245 +246 +247 +248 +249 +250 +251 +252 +253 +254 +255 +256 +257 +258 +259 +260 +261 +262 +263 +264 +265 +266 +267 +268 +269 +270 +271 +272 +273 +274 +275 +276 +277 +278 +279 +280 +281 +282 +283 +284 +285 +286 +287 +288 +289 +290 +291 +292 +293 +294 +295 +296 +297 +298 +299 +300 +301 +302 +303 +304 |
|
MaskedFeedforwardBlock
+
+
+
+ Bases: nn.Module
A feedforward block based on a masked linear module.
+NOTE In this implementation, the number of output features is taken to be equal to the number of input features.
+ + +normflows/nets/made.py
84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 +131 +132 +133 +134 +135 +136 +137 |
|
MaskedLinear
+
+
+
+ Bases: nn.Linear
A linear module with a masked weight matrix.
+ + +normflows/nets/made.py
19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 +77 +78 +79 +80 +81 |
|
MaskedResidualBlock
+
+
+
+ Bases: nn.Module
A residual block containing masked linear modules.
+ + +normflows/nets/made.py
140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 +179 +180 +181 +182 +183 +184 +185 +186 +187 +188 +189 +190 +191 +192 +193 +194 +195 +196 +197 +198 +199 +200 +201 +202 +203 +204 +205 +206 +207 +208 +209 +210 +211 +212 +213 +214 |
|
made_test
+
+
+Tests for MADE. +Code partially taken from https://github.com/bayesiains/nsf
+ + + +mlp
+
+
+MLP
+
+
+
+ Bases: nn.Module
A multilayer perceptron with Leaky ReLU nonlinearities
+ + +normflows/nets/mlp.py
5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 |
|
__init__(layers, leaky=0.0, score_scale=None, output_fn=None, output_scale=None, init_zeros=False, dropout=None)
+
+layers: list of layer sizes from start to end
+leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used
+score_scale: Factor to apply to the scores, i.e. output before output_fn.
+output_fn: String, function to be applied to the output, either None, "sigmoid", "relu", "tanh", or "clampexp"
+output_scale: Rescale outputs if output_fn is specified, i.e. scale * output_fn(out / scale)
+init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see arXiv 1807.03039)
+dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done
normflows/nets/mlp.py
10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 |
|
resnet
+
+
+ResidualBlock
+
+
+
+ Bases: nn.Module
A general-purpose residual block. Works only with 1-dim inputs.
+ + +normflows/nets/resnet.py
7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 |
|
ResidualNet
+
+
+
+ Bases: nn.Module
A general-purpose residual network. Works only with 1-dim inputs.
+ + +normflows/nets/resnet.py
53 + 54 + 55 + 56 + 57 + 58 + 59 + 60 + 61 + 62 + 63 + 64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 |
|
sampling
+
+
+hais
+
+
+HAIS
+
+
+Class which performs HAIS
+ + +normflows/sampling/hais.py
8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 |
|
__init__(betas, prior, target, num_leapfrog, step_size, log_mass)
+
+Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
betas |
+ + | +Annealing schedule, the jth target is |
+ + required + | +
prior |
+ + | +The prior distribution to start the HAIS chain. |
+ + required + | +
target |
+ + | +The target distribution from which we would like to draw weighted samples. |
+ + required + | +
num_leapfrog |
+ + | +Number of leapfrog steps in the HMC transitions. |
+ + required + | +
step_size |
+ + | +step_size to use for HMC transitions. |
+ + required + | +
log_mass |
+ + | +log_mass to use for HMC transitions. |
+ + required + | +
normflows/sampling/hais.py
13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 |
|
sample(num_samples)
+
+Run HAIS to draw samples from the target with appropriate weights.
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
num_samples |
+ + | +The number of samples to draw.a |
+ + required + | +
normflows/sampling/hais.py
37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 |
|
transforms
+
+
+Logit
+
+
+Logit mapping of image tensor, see RealNVP paper
+logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x))
+
+
+
+ normflows/transforms.py
8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 |
|
__init__(alpha=0.05)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
alpha |
+ + | +Alpha parameter, see above |
+
+ 0.05
+ |
+
normflows/transforms.py
17 +18 +19 +20 +21 +22 +23 +24 |
|
Shift
+
+
+Shift data by a fixed constant
+Default is -0.5 to shift data from +interval [0, 1] to [-0.5, 0.5]
+ + +normflows/transforms.py
50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 +64 +65 +66 +67 +68 +69 +70 +71 +72 +73 +74 +75 +76 |
|
__init__(shift=-0.5)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shift |
+ + | +Shift to apply to the data |
+
+ -0.5
+ |
+
normflows/transforms.py
57 +58 +59 +60 +61 +62 +63 +64 |
|
transforms_test
+
+
+utils
+
+
+eval
+
+
+bitsPerDim(model, x, y=None, trans='logit', trans_param=[0.05])
+
+Computes the bits per dim for a batch of data
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
model |
+ + | +Model to compute bits per dim for |
+ + required + | +
x |
+ + | +Batch of data |
+ + required + | +
y |
+ + | +Class labels for batch of data if base distribution is class conditional |
+
+ None
+ |
+
trans |
+ + | +Transformation to be applied to images during training |
+
+ 'logit'
+ |
+
trans_param |
+ + | +List of parameters of the transformation |
+
+ [0.05]
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Bits per dim for data batch under model |
+
normflows/utils/eval.py
5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 +33 +34 |
|
bitsPerDimDataset(model, data_loader, class_cond=True, trans='logit', trans_param=[0.05])
+
+Computes average bits per dim for an entire dataset given by a data loader
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
model |
+ + | +Model to compute bits per dim for |
+ + required + | +
data_loader |
+ + | +Data loader of dataset |
+ + required + | +
class_cond |
+ + | +Flag indicating whether model is class_conditional |
+
+ True
+ |
+
trans |
+ + | +Transformation to be applied to images during training |
+
+ 'logit'
+ |
+
trans_param |
+ + | +List of parameters of the transformation |
+
+ [0.05]
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Average bits per dim for dataset |
+
normflows/utils/eval.py
37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 +62 +63 |
|
masks
+
+
+create_alternating_binary_mask(features, even=True)
+
+Creates a binary mask of a given dimension which alternates its masking.
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
features |
+ + | +Dimension of mask. |
+ + required + | +
even |
+ + | +If True, even values are assigned 1s, odd 0s. If False, vice versa. |
+
+ True
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Alternating binary mask of type torch.Tensor. |
+
normflows/utils/masks.py
4 + 5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 |
|
create_mid_split_binary_mask(features)
+
+Creates a binary mask of a given dimension which splits its masking at the midpoint.
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
features |
+ + | +Dimension of mask. |
+ + required + | +
Returns:
+Type | +Description | +
---|---|
+ | +Binary mask split at midpoint of type torch.Tensor |
+
normflows/utils/masks.py
20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +30 +31 +32 |
|
create_random_binary_mask(features, seed=None)
+
+Creates a random binary mask of a given dimension with half of its entries randomly set to 1s.
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
features |
+ + | +Dimension of mask. |
+ + required + | +
seed |
+ + | +Seed to be used |
+
+ None
+ |
+
Returns:
+Type | +Description | +
---|---|
+ | +Binary mask with half of its entries set to 1s, of type torch.Tensor. |
+
normflows/utils/masks.py
35 +36 +37 +38 +39 +40 +41 +42 +43 +44 +45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 |
|
nn
+
+
+ActNorm
+
+
+
+ Bases: nn.Module
ActNorm layer with just one forward pass
+ + +normflows/utils/nn.py
26 +27 +28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 +43 |
|
__init__(shape)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
shape |
+ + | +Same as shape in flows.ActNorm |
+ + required + | +
logscale_factor |
+ + | +Same as shape in flows.ActNorm |
+ + required + | +
normflows/utils/nn.py
30 +31 +32 +33 +34 +35 +36 +37 +38 +39 |
|
ClampExp
+
+
+
+ Bases: nn.Module
Nonlinearity min(exp(lam * x), 1)
+ + +normflows/utils/nn.py
46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 +58 +59 +60 +61 |
|
__init__()
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
lam |
+ + | +Lambda parameter |
+ + required + | +
normflows/utils/nn.py
51 +52 +53 +54 +55 +56 +57 |
|
ConstScaleLayer
+
+
+
+ Bases: nn.Module
Scaling features by a fixed factor
+ + +normflows/utils/nn.py
7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 |
|
__init__(scale=1.0)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
scale |
+ + | +Scale to apply to features |
+
+ 1.0
+ |
+
normflows/utils/nn.py
12 +13 +14 +15 +16 +17 +18 +19 +20 |
|
PeriodicFeaturesCat
+
+
+
+ Bases: nn.Module
Converts a specified part of the input to periodic features by +replacing those features f with [sin(scale * f), cos(scale * f)].
+Note that this decreases the number of features and their order +is changed.
+ + +normflows/utils/nn.py
133 +134 +135 +136 +137 +138 +139 +140 +141 +142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 +170 +171 +172 +173 +174 +175 +176 +177 +178 |
|
__init__(ndim, ind, scale=1.0)
+
+Constructor +:param ndim: Int, number of dimensions +:param ind: Iterable, indices of input elements to convert to +periodic features +:param scale: Scalar or iterable, used to scale inputs before +converting them to periodic features
+ +normflows/utils/nn.py
142 +143 +144 +145 +146 +147 +148 +149 +150 +151 +152 +153 +154 +155 +156 +157 +158 +159 +160 +161 +162 +163 +164 +165 +166 +167 +168 +169 |
|
PeriodicFeaturesElementwise
+
+
+
+ Bases: nn.Module
Converts a specified part of the input to periodic features by +replacing those features f with +w1 * sin(scale * f) + w2 * cos(scale * f).
+Note that this operation is done elementwise and, therefore, +some information about the feature can be lost.
+ + +normflows/utils/nn.py
64 + 65 + 66 + 67 + 68 + 69 + 70 + 71 + 72 + 73 + 74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 +119 +120 +121 +122 +123 +124 +125 +126 +127 +128 +129 +130 |
|
__init__(ndim, ind, scale=1.0, bias=False, activation=None)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
ndim |
+
+ int
+ |
+ number of dimensions |
+ + required + | +
ind |
+
+ iterable
+ |
+ indices of input elements to convert to periodic features |
+ + required + | +
scale |
+ + | +Scalar or iterable, used to scale inputs before converting them to periodic features |
+
+ 1.0
+ |
+
bias |
+ + | +Flag, whether to add a bias |
+
+ False
+ |
+
activation |
+ + | +Function or None, activation function to be applied |
+
+ None
+ |
+
normflows/utils/nn.py
74 + 75 + 76 + 77 + 78 + 79 + 80 + 81 + 82 + 83 + 84 + 85 + 86 + 87 + 88 + 89 + 90 + 91 + 92 + 93 + 94 + 95 + 96 + 97 + 98 + 99 +100 +101 +102 +103 +104 +105 +106 +107 +108 +109 +110 +111 +112 +113 +114 +115 +116 +117 +118 |
|
sum_except_batch(x, num_batch_dims=1)
+
+Sums all elements of x
except for the first num_batch_dims
dimensions.
normflows/utils/nn.py
190 +191 +192 +193 |
|
optim
+
+
+clear_grad(model)
+
+Set gradients of model parameter to None as this speeds up training,
+See youtube
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
model |
+ + | +Model to clear gradients of |
+ + required + | +
normflows/utils/optim.py
16 +17 +18 +19 +20 +21 +22 +23 +24 +25 |
|
set_requires_grad(module, flag)
+
+Sets requires_grad flag of all parameters of a torch.nn.module
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
module |
+ + | +torch.nn.module |
+ + required + | +
flag |
+ + | +Flag to set requires_grad to |
+ + required + | +
normflows/utils/optim.py
4 + 5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 |
|
preprocessing
+
+
+Jitter
+
+
+Transform for dataloader, adds uniform jitter noise to data
+ + +normflows/utils/preprocessing.py
28 +29 +30 +31 +32 +33 +34 +35 +36 +37 +38 +39 +40 +41 +42 |
|
__init__(scale=1.0 / 256)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
scale |
+ + | +Scaling factor for noise |
+
+ 1.0 / 256
+ |
+
normflows/utils/preprocessing.py
31 +32 +33 +34 +35 +36 +37 |
|
Logit
+
+
+Transform for dataloader
+logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x))
+
+
+
+ normflows/utils/preprocessing.py
4 + 5 + 6 + 7 + 8 + 9 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +21 +22 +23 +24 +25 |
|
__init__(alpha=0)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
alpha |
+ + | +see above |
+
+ 0
+ |
+
normflows/utils/preprocessing.py
12 +13 +14 +15 +16 +17 +18 |
|
Scale
+
+
+Transform for dataloader, adds uniform jitter noise to data
+ + +normflows/utils/preprocessing.py
45 +46 +47 +48 +49 +50 +51 +52 +53 +54 +55 +56 +57 |
|
__init__(scale=255.0 / 256.0)
+
+Constructor
+ +Parameters:
+Name | +Type | +Description | +Default | +
---|---|---|---|
scale |
+ + | +Scaling factor for noise |
+
+ 255.0 / 256.0
+ |
+
normflows/utils/preprocessing.py
48 +49 +50 +51 +52 +53 +54 |
|
' + escapeHtml(summary) +'
' + noResultsText + '
'); + } +} + +function doSearch () { + var query = document.getElementById('mkdocs-search-query').value; + if (query.length > min_search_length) { + if (!window.Worker) { + displayResults(search(query)); + } else { + searchWorker.postMessage({query: query}); + } + } else { + // Clear results for short queries + displayResults([]); + } +} + +function initSearch () { + var search_input = document.getElementById('mkdocs-search-query'); + if (search_input) { + search_input.addEventListener("keyup", doSearch); + } + var term = getSearchTermFromLocation(); + if (term) { + search_input.value = term; + doSearch(); + } +} + +function onWorkerMessage (e) { + if (e.data.allowSearch) { + initSearch(); + } else if (e.data.results) { + var results = e.data.results; + displayResults(results); + } else if (e.data.config) { + min_search_length = e.data.config.min_search_length-1; + } +} + +if (!window.Worker) { + console.log('Web Worker API not supported'); + // load index in main thread + $.getScript(joinUrl(base_url, "search/worker.js")).done(function () { + console.log('Loaded worker'); + init(); + window.postMessage = function (msg) { + onWorkerMessage({data: msg}); + }; + }).fail(function (jqxhr, settings, exception) { + console.error('Could not load worker.js'); + }); +} else { + // Wrap search in a web worker + var searchWorker = new Worker(joinUrl(base_url, "search/worker.js")); + searchWorker.postMessage({init: true}); + searchWorker.onmessage = onWorkerMessage; +} diff --git a/search/search_index.json b/search/search_index.json new file mode 100644 index 0000000..8a3895e --- /dev/null +++ b/search/search_index.json @@ -0,0 +1 @@ +{"config":{"indexing":"full","lang":["en"],"min_search_length":3,"prebuild_index":false,"separator":"[\\s\\-]+"},"docs":[{"location":"","text":"normflows : A PyTorch Package for Normalizing Flows normflows is a PyTorch implementation of discrete normalizing flows. Many popular flow architectures are implemented, see the list below . The package can be easily installed via pip . The basic usage is described here , and a full documentation is available as well. A more detailed description of this package is given in out accompanying paper . Several sample use cases are provided in the examples folder , including Glow , a VAE , and a Residual Flow . Moreover, two simple applications are highlighed in the examples section . You can run them yourself in Google Colab using the links below to get a feeling for normflows . Link Description Real NVP applied to a 2D bimodal target distribution Modeling a distribution on a cylinder surface with a neural spline flow Modeling and generating CIFAR-10 images with Glow Implemented Flows Architecture Reference Planar Flow Rezende & Mohamed, 2015 Radial Flow Rezende & Mohamed, 2015 NICE Dinh et al., 2014 Real NVP Dinh et al., 2017 Glow Kingma et al., 2018 Masked Autoregressive Flow Papamakarios et al., 2017 Neural Spline Flow Durkan et al., 2019 Circular Neural Spline Flow Rezende et al., 2020 Residual Flow Chen et al., 2019 Stochastic Normalizing Flow Wu et al., 2020 Note that Neural Spline Flows with circular and non-circular coordinates are supported as well. Installation The latest version of the package can be installed via pip pip install normflows At least Python 3.7 is required. If you want to use a GPU, make sure that PyTorch is set up correctly by following the instructions at the PyTorch website . To run the example notebooks clone the repository first git clone https://github.com/VincentStimper/normalizing-flows.git and then install the dependencies. pip install -r requirements_examples.txt Usage A normalizing flow consists of a base distribution, defined in nf.distributions.base , and a list of flows, given in nf.flows . Let's assume our target is a 2D distribution. We pick a diagonal Gaussian base distribution, which is the most popular choice. Our flow shall be a Real NVP model and, therefore, we need to define a neural network for computing the parameters of the affine coupling map. One dimension is used to compute the scale and shift parameter for the other dimension. After each coupling layer we swap their roles. import normflows as nf # Define 2D Gaussian base distribution base = nf.distributions.base.DiagGaussian(2) # Define list of flows num_layers = 32 flows = [] for i in range(num_layers): # Neural network with two hidden layers having 64 units each # Last layer is initialized by zeros making training more stable param_map = nf.nets.MLP([1, 64, 64, 2], init_zeros=True) # Add flow layer flows.append(nf.flows.AffineCouplingBlock(param_map)) # Swap dimensions flows.append(nf.flows.Permute(2, mode='swap')) Once they are set up, we can define a nf.NormalizingFlow model. If the target density is available, it can be added to the model to be used during training. Sample target distributions are given in nf.distributions.target . # If the target density is not given model = nf.NormalizingFlow(base, flows) # If the target density is given target = nf.distributions.target.TwoMoons() model = nf.NormalizingFlow(base, flows, target) The loss can be computed with the methods of the model and minimized. # When doing maximum likelihood learning, i.e. minimizing the forward KLD # with no target distribution given loss = model.forward_kld(x) # When minimizing the reverse KLD based on the given target distribution loss = model.reverse_kld(num_samples=512) # Optimization as usual loss.backward() optimizer.step() Examples We provide several illustrative examples of how to use the package in the examples directory. Amoung them are implementations of Glow , a VAE , and a Residual Flow . More advanced experiments can be done with the scripts listed in the repository about resampled base distributions , see its experiments folder. Below, we consider two simple 2D examples. Real NVP applied to a 2D bimodal target distribution In this notebook , which can directly be opened in Colab , we consider a 2D distribution with two half-moon-shaped modes as a target. We approximate it with a Real NVP model and obtain the following results. Note that there might be a density filament connecting the two modes, which is due to an architectural limitation of normalizing flows, especially prominent in Real NVP. You can find out more about it in this paper . Modeling a distribution on a cylinder surface with a neural spline flow In another example , which is available in Colab as well, we apply a Neural Spline Flow model to a distribution defined on a cylinder. The resulting density is visualized below. This example is considered in the paper accompanying this repository. Support If you have problems, please read the package documentation and check out the examples section above. You are also welcome to create issues on GitHub to get help. Note that it is worthwhile browsing the existing open and closed issues, which might address the problem you are facing. Contributing If you find a bug or have a feature request, please file an issue on GitHub . You are welcome to contribute to the package by fixing the bug or adding the feature yourself. If you want to contribute, please add tests for the code you added or modified and ensure it passes successfully by running pytest . This can be done by simply executing pytest within your local version of the repository. Make sure you code is well documented, and we also encourage contributions to the existing documentation. Once you finished coding and testing, please create a pull request on GitHub . Used by The package has been used in several research papers, which are listed below. Andrew Campbell, Wenlong Chen, Vincent Stimper, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato, and Yichuan Zhang. A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization . In Proceedings of the 38th International Conference on Machine Learning, pp. 1238\u20131248. PMLR, 2021. Code available on GitHub. Vincent Stimper, Bernhard Sch\u00f6lkopf, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato. Resampling Base Distributions of Normalizing Flows . In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 4915\u20134936, 2022. Code available on GitHub. Laurence I. Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Sch\u00f6lkopf, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato. Flow Annealed Importance Sampling Bootstrap . arXiv preprint arXiv:2208.01893, 2022. Code available on GitHub. Moreover, the boltzgen package has been build upon normflows . Citation If you use normflows , please consider citing the corresponding paper as follows. Vincent Stimper, David Liu, Andrew Campbell, Vincent Berenz, Lukas Ryll, Bernhard Sch\u00f6lkopf, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato. normflows: A PyTorch Package for Normalizing Flows, arXiv preprint arXiv:2302.12014, 2023. Bibtex @article{normflows, author = {Vincent Stimper and David Liu and Andrew Campbell and Vincent Berenz and Lukas Ryll and Bernhard Sch{\\\"o}lkopf and Jos{\\'e} Miguel Hern{\\'a}ndez-Lobato}, title = {normflows: {A} {P}y{T}orch {P}ackage for {N}ormalizing {F}lows}, journal = {arXiv preprint arXiv:2302.12014}, year = {2023} }","title":"about"},{"location":"#normflows-a-pytorch-package-for-normalizing-flows","text":"normflows is a PyTorch implementation of discrete normalizing flows. Many popular flow architectures are implemented, see the list below . The package can be easily installed via pip . The basic usage is described here , and a full documentation is available as well. A more detailed description of this package is given in out accompanying paper . Several sample use cases are provided in the examples folder , including Glow , a VAE , and a Residual Flow . Moreover, two simple applications are highlighed in the examples section . You can run them yourself in Google Colab using the links below to get a feeling for normflows . Link Description Real NVP applied to a 2D bimodal target distribution Modeling a distribution on a cylinder surface with a neural spline flow Modeling and generating CIFAR-10 images with Glow","title":"normflows: A PyTorch Package for Normalizing Flows"},{"location":"#implemented-flows","text":"Architecture Reference Planar Flow Rezende & Mohamed, 2015 Radial Flow Rezende & Mohamed, 2015 NICE Dinh et al., 2014 Real NVP Dinh et al., 2017 Glow Kingma et al., 2018 Masked Autoregressive Flow Papamakarios et al., 2017 Neural Spline Flow Durkan et al., 2019 Circular Neural Spline Flow Rezende et al., 2020 Residual Flow Chen et al., 2019 Stochastic Normalizing Flow Wu et al., 2020 Note that Neural Spline Flows with circular and non-circular coordinates are supported as well.","title":"Implemented Flows"},{"location":"#installation","text":"The latest version of the package can be installed via pip pip install normflows At least Python 3.7 is required. If you want to use a GPU, make sure that PyTorch is set up correctly by following the instructions at the PyTorch website . To run the example notebooks clone the repository first git clone https://github.com/VincentStimper/normalizing-flows.git and then install the dependencies. pip install -r requirements_examples.txt","title":"Installation"},{"location":"#usage","text":"A normalizing flow consists of a base distribution, defined in nf.distributions.base , and a list of flows, given in nf.flows . Let's assume our target is a 2D distribution. We pick a diagonal Gaussian base distribution, which is the most popular choice. Our flow shall be a Real NVP model and, therefore, we need to define a neural network for computing the parameters of the affine coupling map. One dimension is used to compute the scale and shift parameter for the other dimension. After each coupling layer we swap their roles. import normflows as nf # Define 2D Gaussian base distribution base = nf.distributions.base.DiagGaussian(2) # Define list of flows num_layers = 32 flows = [] for i in range(num_layers): # Neural network with two hidden layers having 64 units each # Last layer is initialized by zeros making training more stable param_map = nf.nets.MLP([1, 64, 64, 2], init_zeros=True) # Add flow layer flows.append(nf.flows.AffineCouplingBlock(param_map)) # Swap dimensions flows.append(nf.flows.Permute(2, mode='swap')) Once they are set up, we can define a nf.NormalizingFlow model. If the target density is available, it can be added to the model to be used during training. Sample target distributions are given in nf.distributions.target . # If the target density is not given model = nf.NormalizingFlow(base, flows) # If the target density is given target = nf.distributions.target.TwoMoons() model = nf.NormalizingFlow(base, flows, target) The loss can be computed with the methods of the model and minimized. # When doing maximum likelihood learning, i.e. minimizing the forward KLD # with no target distribution given loss = model.forward_kld(x) # When minimizing the reverse KLD based on the given target distribution loss = model.reverse_kld(num_samples=512) # Optimization as usual loss.backward() optimizer.step()","title":"Usage"},{"location":"#examples","text":"We provide several illustrative examples of how to use the package in the examples directory. Amoung them are implementations of Glow , a VAE , and a Residual Flow . More advanced experiments can be done with the scripts listed in the repository about resampled base distributions , see its experiments folder. Below, we consider two simple 2D examples.","title":"Examples"},{"location":"#real-nvp-applied-to-a-2d-bimodal-target-distribution","text":"In this notebook , which can directly be opened in Colab , we consider a 2D distribution with two half-moon-shaped modes as a target. We approximate it with a Real NVP model and obtain the following results. Note that there might be a density filament connecting the two modes, which is due to an architectural limitation of normalizing flows, especially prominent in Real NVP. You can find out more about it in this paper .","title":"Real NVP applied to a 2D bimodal target distribution"},{"location":"#modeling-a-distribution-on-a-cylinder-surface-with-a-neural-spline-flow","text":"In another example , which is available in Colab as well, we apply a Neural Spline Flow model to a distribution defined on a cylinder. The resulting density is visualized below. This example is considered in the paper accompanying this repository.","title":"Modeling a distribution on a cylinder surface with a neural spline flow"},{"location":"#support","text":"If you have problems, please read the package documentation and check out the examples section above. You are also welcome to create issues on GitHub to get help. Note that it is worthwhile browsing the existing open and closed issues, which might address the problem you are facing.","title":"Support"},{"location":"#contributing","text":"If you find a bug or have a feature request, please file an issue on GitHub . You are welcome to contribute to the package by fixing the bug or adding the feature yourself. If you want to contribute, please add tests for the code you added or modified and ensure it passes successfully by running pytest . This can be done by simply executing pytest within your local version of the repository. Make sure you code is well documented, and we also encourage contributions to the existing documentation. Once you finished coding and testing, please create a pull request on GitHub .","title":"Contributing"},{"location":"#used-by","text":"The package has been used in several research papers, which are listed below. Andrew Campbell, Wenlong Chen, Vincent Stimper, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato, and Yichuan Zhang. A gradient based strategy for Hamiltonian Monte Carlo hyperparameter optimization . In Proceedings of the 38th International Conference on Machine Learning, pp. 1238\u20131248. PMLR, 2021. Code available on GitHub. Vincent Stimper, Bernhard Sch\u00f6lkopf, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato. Resampling Base Distributions of Normalizing Flows . In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151, pp. 4915\u20134936, 2022. Code available on GitHub. Laurence I. Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Sch\u00f6lkopf, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato. Flow Annealed Importance Sampling Bootstrap . arXiv preprint arXiv:2208.01893, 2022. Code available on GitHub. Moreover, the boltzgen package has been build upon normflows .","title":"Used by"},{"location":"#citation","text":"If you use normflows , please consider citing the corresponding paper as follows. Vincent Stimper, David Liu, Andrew Campbell, Vincent Berenz, Lukas Ryll, Bernhard Sch\u00f6lkopf, Jos\u00e9 Miguel Hern\u00e1ndez-Lobato. normflows: A PyTorch Package for Normalizing Flows, arXiv preprint arXiv:2302.12014, 2023. Bibtex @article{normflows, author = {Vincent Stimper and David Liu and Andrew Campbell and Vincent Berenz and Lukas Ryll and Bernhard Sch{\\\"o}lkopf and Jos{\\'e} Miguel Hern{\\'a}ndez-Lobato}, title = {normflows: {A} {P}y{T}orch {P}ackage for {N}ormalizing {F}lows}, journal = {arXiv preprint arXiv:2302.12014}, year = {2023} }","title":"Citation"},{"location":"references/","text":"API references core ClassCondFlow Bases: nn . Module Class conditional normalizing Flow model Source code in normflows/core.py 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 class ClassCondFlow ( nn . Module ): \"\"\" Class conditional normalizing Flow model \"\"\" def __init__ ( self , q0 , flows ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) def forward_kld ( self , x , y ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return - torch . mean ( log_q ) def sample ( self , num_samples = 1 , y = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples , y ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return log_q def save ( self , path ): \"\"\"Save state dict of model Args: param path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) __init__ ( q0 , flows ) Constructor Parameters: Name Type Description Default q0 Base distribution required flows List of flows required Source code in normflows/core.py 221 222 223 224 225 226 227 228 229 230 def __init__ ( self , q0 , flows ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) forward_kld ( x , y ) Estimates forward KL divergence, see arXiv 1912.02762 Parameters: Name Type Description Default x Batch sampled from target distribution required Returns: Type Description Estimate of forward KL divergence averaged over batch Source code in normflows/core.py 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 def forward_kld ( self , x , y ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return - torch . mean ( log_q ) load ( path ) Load model from state dict Parameters: Name Type Description Default path Path including filename where to load model from required Source code in normflows/core.py 291 292 293 294 295 296 297 def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) log_prob ( x , y ) Get log probability for batch Parameters: Name Type Description Default x Batch required y Classes of x required Returns: Type Description log probability Source code in normflows/core.py 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return log_q sample ( num_samples = 1 , y = None ) Samples from flow-based approximate distribution Parameters: Name Type Description Default num_samples Number of samples to draw 1 y Classes to sample from, will be sampled uniformly if None None Returns: Type Description Samples, log probability Source code in normflows/core.py 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 def sample ( self , num_samples = 1 , y = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples , y ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q save ( path ) Save state dict of model Parameters: Name Type Description Default param path Path including filename where to save model required Source code in normflows/core.py 283 284 285 286 287 288 289 def save ( self , path ): \"\"\"Save state dict of model Args: param path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) MultiscaleFlow Bases: nn . Module Normalizing Flow model with multiscale architecture, see RealNVP or Glow paper Source code in normflows/core.py 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 class MultiscaleFlow ( nn . Module ): \"\"\" Normalizing Flow model with multiscale architecture, see RealNVP or Glow paper \"\"\" def __init__ ( self , q0 , flows , merges , transform = None , class_cond = True ): \"\"\"Constructor Args: q0: List of base distribution flows: List of list of flows for each level merges: List of merge/split operations (forward pass must do merge) transform: Initial transformation of inputs class_cond: Flag, indicated whether model has class conditional base distributions \"\"\" super () . __init__ () self . q0 = nn . ModuleList ( q0 ) self . num_levels = len ( self . q0 ) self . flows = torch . nn . ModuleList ([ nn . ModuleList ( flow ) for flow in flows ]) self . merges = torch . nn . ModuleList ( merges ) self . transform = transform self . class_cond = class_cond def forward_kld ( self , x , y = None ): \"\"\"Estimates forward KL divergence, see see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution y: Batch of targets, if applicable Returns: Estimate of forward KL divergence averaged over batch \"\"\" return - torch . mean ( self . log_prob ( x , y )) def forward ( self , x , y = None ): \"\"\"Get negative log-likelihood for maximum likelihood training Args: x: Batch of data y: Batch of targets, if applicable Returns: Negative log-likelihood of the batch \"\"\" return - self . log_prob ( x , y ) def sample ( self , num_samples = 1 , y = None , temperature = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None temperature: Temperature parameter for temp annealed sampling Returns: Samples, log probability \"\"\" if temperature is not None : self . set_temperature ( temperature ) for i in range ( len ( self . q0 )): if self . class_cond : z_ , log_q_ = self . q0 [ i ]( num_samples , y ) else : z_ , log_q_ = self . q0 [ i ]( num_samples ) if i == 0 : log_q = log_q_ z = z_ else : log_q += log_q_ z , log_det = self . merges [ i - 1 ]([ z , z_ ]) log_q -= log_det for flow in self . flows [ i ]: z , log_det = flow ( z ) log_q -= log_det if self . transform is not None : z , log_det = self . transform ( z ) log_q -= log_det if temperature is not None : self . reset_temperature () return z , log_q def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = 0 z = x if self . transform is not None : z , log_det = self . transform . inverse ( z ) log_q += log_det for i in range ( len ( self . q0 ) - 1 , - 1 , - 1 ): for j in range ( len ( self . flows [ i ]) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ][ j ] . inverse ( z ) log_q += log_det if i > 0 : [ z , z_ ], log_det = self . merges [ i - 1 ] . inverse ( z ) log_q += log_det else : z_ = z if self . class_cond : log_q += self . q0 [ i ] . log_prob ( z_ , y ) else : log_q += self . q0 [ i ] . log_prob ( z_ ) return log_q def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) def set_temperature ( self , temperature ): \"\"\"Set temperature for temperature a annealed sampling Args: temperature: Temperature parameter \"\"\" for q0 in self . q0 : if hasattr ( q0 , \"temperature\" ): q0 . temperature = temperature else : raise NotImplementedError ( \"One base function does not \" \"support temperature annealed sampling\" ) def reset_temperature ( self ): \"\"\" Set temperature values of base distributions back to None \"\"\" self . set_temperature ( None ) __init__ ( q0 , flows , merges , transform = None , class_cond = True ) Constructor Parameters: Name Type Description Default q0 List of base distribution required flows List of list of flows for each level required merges List of merge/split operations (forward pass must do merge) required transform Initial transformation of inputs None class_cond Flag, indicated whether model has class conditional True base distributions Source code in normflows/core.py 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 def __init__ ( self , q0 , flows , merges , transform = None , class_cond = True ): \"\"\"Constructor Args: q0: List of base distribution flows: List of list of flows for each level merges: List of merge/split operations (forward pass must do merge) transform: Initial transformation of inputs class_cond: Flag, indicated whether model has class conditional base distributions \"\"\" super () . __init__ () self . q0 = nn . ModuleList ( q0 ) self . num_levels = len ( self . q0 ) self . flows = torch . nn . ModuleList ([ nn . ModuleList ( flow ) for flow in flows ]) self . merges = torch . nn . ModuleList ( merges ) self . transform = transform self . class_cond = class_cond forward ( x , y = None ) Get negative log-likelihood for maximum likelihood training Parameters: Name Type Description Default x Batch of data required y Batch of targets, if applicable None Returns: Type Description Negative log-likelihood of the batch Source code in normflows/core.py 337 338 339 340 341 342 343 344 345 346 347 def forward ( self , x , y = None ): \"\"\"Get negative log-likelihood for maximum likelihood training Args: x: Batch of data y: Batch of targets, if applicable Returns: Negative log-likelihood of the batch \"\"\" return - self . log_prob ( x , y ) forward_kld ( x , y = None ) Estimates forward KL divergence, see see arXiv 1912.02762 Parameters: Name Type Description Default x Batch sampled from target distribution required y Batch of targets, if applicable None Returns: Type Description Estimate of forward KL divergence averaged over batch Source code in normflows/core.py 325 326 327 328 329 330 331 332 333 334 335 def forward_kld ( self , x , y = None ): \"\"\"Estimates forward KL divergence, see see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution y: Batch of targets, if applicable Returns: Estimate of forward KL divergence averaged over batch \"\"\" return - torch . mean ( self . log_prob ( x , y )) load ( path ) Load model from state dict Parameters: Name Type Description Default path Path including filename where to load model from required Source code in normflows/core.py 422 423 424 425 426 427 428 def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) log_prob ( x , y ) Get log probability for batch Parameters: Name Type Description Default x Batch required y Classes of x required Returns: Type Description log probability Source code in normflows/core.py 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = 0 z = x if self . transform is not None : z , log_det = self . transform . inverse ( z ) log_q += log_det for i in range ( len ( self . q0 ) - 1 , - 1 , - 1 ): for j in range ( len ( self . flows [ i ]) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ][ j ] . inverse ( z ) log_q += log_det if i > 0 : [ z , z_ ], log_det = self . merges [ i - 1 ] . inverse ( z ) log_q += log_det else : z_ = z if self . class_cond : log_q += self . q0 [ i ] . log_prob ( z_ , y ) else : log_q += self . q0 [ i ] . log_prob ( z_ ) return log_q reset_temperature () Set temperature values of base distributions back to None Source code in normflows/core.py 445 446 447 448 449 def reset_temperature ( self ): \"\"\" Set temperature values of base distributions back to None \"\"\" self . set_temperature ( None ) sample ( num_samples = 1 , y = None , temperature = None ) Samples from flow-based approximate distribution Parameters: Name Type Description Default num_samples Number of samples to draw 1 y Classes to sample from, will be sampled uniformly if None None temperature Temperature parameter for temp annealed sampling None Returns: Type Description Samples, log probability Source code in normflows/core.py 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 def sample ( self , num_samples = 1 , y = None , temperature = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None temperature: Temperature parameter for temp annealed sampling Returns: Samples, log probability \"\"\" if temperature is not None : self . set_temperature ( temperature ) for i in range ( len ( self . q0 )): if self . class_cond : z_ , log_q_ = self . q0 [ i ]( num_samples , y ) else : z_ , log_q_ = self . q0 [ i ]( num_samples ) if i == 0 : log_q = log_q_ z = z_ else : log_q += log_q_ z , log_det = self . merges [ i - 1 ]([ z , z_ ]) log_q -= log_det for flow in self . flows [ i ]: z , log_det = flow ( z ) log_q -= log_det if self . transform is not None : z , log_det = self . transform ( z ) log_q -= log_det if temperature is not None : self . reset_temperature () return z , log_q save ( path ) Save state dict of model Parameters: Name Type Description Default path Path including filename where to save model required Source code in normflows/core.py 414 415 416 417 418 419 420 def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) set_temperature ( temperature ) Set temperature for temperature a annealed sampling Parameters: Name Type Description Default temperature Temperature parameter required Source code in normflows/core.py 430 431 432 433 434 435 436 437 438 439 440 441 442 443 def set_temperature ( self , temperature ): \"\"\"Set temperature for temperature a annealed sampling Args: temperature: Temperature parameter \"\"\" for q0 in self . q0 : if hasattr ( q0 , \"temperature\" ): q0 . temperature = temperature else : raise NotImplementedError ( \"One base function does not \" \"support temperature annealed sampling\" ) NormalizingFlow Bases: nn . Module Normalizing Flow model to approximate target distribution Source code in normflows/core.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 class NormalizingFlow ( nn . Module ): \"\"\" Normalizing Flow model to approximate target distribution \"\"\" def __init__ ( self , q0 , flows , p = None ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows p: Target distribution \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) self . p = p def forward ( self , z ): \"\"\"Transforms latent variable z to the flow variable x Args: z: Batch in the latent space Returns: Batch in the space of the target distribution \"\"\" for flow in self . flows : z , _ = flow ( z ) return z def forward_and_log_det ( self , z ): \"\"\"Transforms latent variable z to the flow variable x and computes log determinant of the Jacobian Args: z: Batch in the latent space Returns: Batch in the space of the target distribution, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( z ), device = z . device ) for flow in self . flows : z , log_d = flow ( z ) log_det += log_d return z , log_det def inverse ( self , x ): \"\"\"Transforms flow variable x to the latent variable z Args: x: Batch in the space of the target distribution Returns: Batch in the latent space \"\"\" for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , _ = self . flows [ i ] . inverse ( x ) return x def inverse_and_log_det ( self , x ): \"\"\"Transforms flow variable x to the latent variable z and computes log determinant of the Jacobian Args: x: Batch in the space of the target distribution Returns: Batch in the latent space, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( x ), device = x . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , log_d = self . flows [ i ] . inverse ( x ) log_det += log_d return x , log_det def forward_kld ( self , x ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return - torch . mean ( log_q ) def reverse_kld ( self , num_samples = 1 , beta = 1.0 , score_fn = True ): \"\"\"Estimates reverse KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: num_samples: Number of samples to draw from base distribution beta: Annealing parameter, see [arXiv 1505.05770](https://arxiv.org/abs/1505.05770) score_fn: Flag whether to include score function in gradient, see [arXiv 1703.09194](https://arxiv.org/abs/1703.09194) Returns: Estimate of the reverse KL divergence averaged over latent samples \"\"\" z , log_q_ = self . q0 ( num_samples ) log_q = torch . zeros_like ( log_q_ ) log_q += log_q_ for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det if not score_fn : z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) log_p = self . p . log_prob ( z ) return torch . mean ( log_q ) - beta * torch . mean ( log_p ) def reverse_alpha_div ( self , num_samples = 1 , alpha = 1 , dreg = False ): \"\"\"Alpha divergence when sampling from q Args: num_samples: Number of samples to draw dreg: Flag whether to use Double Reparametrized Gradient estimator, see [arXiv 1810.04152](https://arxiv.org/abs/1810.04152) Returns: Alpha divergence \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . p . log_prob ( z ) if dreg : w_const = torch . exp ( log_p - log_q ) . detach () z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) w = torch . exp ( log_p - log_q ) w_alpha = w_const ** alpha w_alpha = w_alpha / torch . mean ( w_alpha ) weights = ( 1 - alpha ) * w_alpha + alpha * w_alpha ** 2 loss = - alpha * torch . mean ( weights * torch . log ( w )) else : loss = np . sign ( alpha - 1 ) * torch . logsumexp ( alpha * ( log_p - log_q ), 0 ) return loss def sample ( self , num_samples = 1 ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q def log_prob ( self , x ): \"\"\"Get log probability for batch Args: x: Batch Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return log_q def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) __init__ ( q0 , flows , p = None ) Constructor Parameters: Name Type Description Default q0 Base distribution required flows List of flows required p Target distribution None Source code in normflows/core.py 14 15 16 17 18 19 20 21 22 23 24 25 def __init__ ( self , q0 , flows , p = None ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows p: Target distribution \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) self . p = p forward ( z ) Transforms latent variable z to the flow variable x Parameters: Name Type Description Default z Batch in the latent space required Returns: Type Description Batch in the space of the target distribution Source code in normflows/core.py 27 28 29 30 31 32 33 34 35 36 37 38 def forward ( self , z ): \"\"\"Transforms latent variable z to the flow variable x Args: z: Batch in the latent space Returns: Batch in the space of the target distribution \"\"\" for flow in self . flows : z , _ = flow ( z ) return z forward_and_log_det ( z ) Transforms latent variable z to the flow variable x and computes log determinant of the Jacobian Parameters: Name Type Description Default z Batch in the latent space required Returns: Type Description Batch in the space of the target distribution, log determinant of the Jacobian Source code in normflows/core.py 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def forward_and_log_det ( self , z ): \"\"\"Transforms latent variable z to the flow variable x and computes log determinant of the Jacobian Args: z: Batch in the latent space Returns: Batch in the space of the target distribution, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( z ), device = z . device ) for flow in self . flows : z , log_d = flow ( z ) log_det += log_d return z , log_det forward_kld ( x ) Estimates forward KL divergence, see arXiv 1912.02762 Parameters: Name Type Description Default x Batch sampled from target distribution required Returns: Type Description Estimate of forward KL divergence averaged over batch Source code in normflows/core.py 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 def forward_kld ( self , x ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return - torch . mean ( log_q ) inverse ( x ) Transforms flow variable x to the latent variable z Parameters: Name Type Description Default x Batch in the space of the target distribution required Returns: Type Description Batch in the latent space Source code in normflows/core.py 57 58 59 60 61 62 63 64 65 66 67 68 def inverse ( self , x ): \"\"\"Transforms flow variable x to the latent variable z Args: x: Batch in the space of the target distribution Returns: Batch in the latent space \"\"\" for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , _ = self . flows [ i ] . inverse ( x ) return x inverse_and_log_det ( x ) Transforms flow variable x to the latent variable z and computes log determinant of the Jacobian Parameters: Name Type Description Default x Batch in the space of the target distribution required Returns: Type Description Batch in the latent space, log determinant of the Jacobian Source code in normflows/core.py 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 def inverse_and_log_det ( self , x ): \"\"\"Transforms flow variable x to the latent variable z and computes log determinant of the Jacobian Args: x: Batch in the space of the target distribution Returns: Batch in the latent space, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( x ), device = x . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , log_d = self . flows [ i ] . inverse ( x ) log_det += log_d return x , log_det load ( path ) Load model from state dict Parameters: Name Type Description Default path Path including filename where to load model from required Source code in normflows/core.py 207 208 209 210 211 212 213 def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) log_prob ( x ) Get log probability for batch Parameters: Name Type Description Default x Batch required Returns: Type Description log probability Source code in normflows/core.py 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 def log_prob ( self , x ): \"\"\"Get log probability for batch Args: x: Batch Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return log_q reverse_alpha_div ( num_samples = 1 , alpha = 1 , dreg = False ) Alpha divergence when sampling from q Parameters: Name Type Description Default num_samples Number of samples to draw 1 dreg Flag whether to use Double Reparametrized Gradient estimator, see arXiv 1810.04152 False Returns: Type Description Alpha divergence Source code in normflows/core.py 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 def reverse_alpha_div ( self , num_samples = 1 , alpha = 1 , dreg = False ): \"\"\"Alpha divergence when sampling from q Args: num_samples: Number of samples to draw dreg: Flag whether to use Double Reparametrized Gradient estimator, see [arXiv 1810.04152](https://arxiv.org/abs/1810.04152) Returns: Alpha divergence \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . p . log_prob ( z ) if dreg : w_const = torch . exp ( log_p - log_q ) . detach () z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) w = torch . exp ( log_p - log_q ) w_alpha = w_const ** alpha w_alpha = w_alpha / torch . mean ( w_alpha ) weights = ( 1 - alpha ) * w_alpha + alpha * w_alpha ** 2 loss = - alpha * torch . mean ( weights * torch . log ( w )) else : loss = np . sign ( alpha - 1 ) * torch . logsumexp ( alpha * ( log_p - log_q ), 0 ) return loss reverse_kld ( num_samples = 1 , beta = 1.0 , score_fn = True ) Estimates reverse KL divergence, see arXiv 1912.02762 Parameters: Name Type Description Default num_samples Number of samples to draw from base distribution 1 beta Annealing parameter, see arXiv 1505.05770 1.0 score_fn Flag whether to include score function in gradient, see arXiv 1703.09194 True Returns: Type Description Estimate of the reverse KL divergence averaged over latent samples Source code in normflows/core.py 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 def reverse_kld ( self , num_samples = 1 , beta = 1.0 , score_fn = True ): \"\"\"Estimates reverse KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: num_samples: Number of samples to draw from base distribution beta: Annealing parameter, see [arXiv 1505.05770](https://arxiv.org/abs/1505.05770) score_fn: Flag whether to include score function in gradient, see [arXiv 1703.09194](https://arxiv.org/abs/1703.09194) Returns: Estimate of the reverse KL divergence averaged over latent samples \"\"\" z , log_q_ = self . q0 ( num_samples ) log_q = torch . zeros_like ( log_q_ ) log_q += log_q_ for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det if not score_fn : z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) log_p = self . p . log_prob ( z ) return torch . mean ( log_q ) - beta * torch . mean ( log_p ) sample ( num_samples = 1 ) Samples from flow-based approximate distribution Parameters: Name Type Description Default num_samples Number of samples to draw 1 Returns: Type Description Samples, log probability Source code in normflows/core.py 167 168 169 170 171 172 173 174 175 176 177 178 179 180 def sample ( self , num_samples = 1 ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q save ( path ) Save state dict of model Parameters: Name Type Description Default path Path including filename where to save model required Source code in normflows/core.py 199 200 201 202 203 204 205 def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) NormalizingFlowVAE Bases: nn . Module VAE using normalizing flows to express approximate distribution Source code in normflows/core.py 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 class NormalizingFlowVAE ( nn . Module ): \"\"\" VAE using normalizing flows to express approximate distribution \"\"\" def __init__ ( self , prior , q0 = distributions . Dirac (), flows = None , decoder = None ): \"\"\"Constructor of normalizing flow model Args: prior: Prior distribution of te VAE, i.e. Gaussian decoder: Optional decoder flows: Flows to transform output of base encoder q0: Base Encoder \"\"\" super () . __init__ () self . prior = prior self . decoder = decoder self . flows = nn . ModuleList ( flows ) self . q0 = q0 def forward ( self , x , num_samples = 1 ): \"\"\"Takes data batch, samples num_samples for each data point from base distribution Args: x: data batch num_samples: number of samples to draw for each data point Returns: latent variables for each batch and sample, log_q, and log_p \"\"\" z , log_q = self . q0 ( x , num_samples = num_samples ) # Flatten batch and sample dim z = z . view ( - 1 , * z . size ()[ 2 :]) log_q = log_q . view ( - 1 , * log_q . size ()[ 2 :]) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . prior . log_prob ( z ) if self . decoder is not None : log_p += self . decoder . log_prob ( x , z ) # Separate batch and sample dimension again z = z . view ( - 1 , num_samples , * z . size ()[ 1 :]) log_q = log_q . view ( - 1 , num_samples , * log_q . size ()[ 1 :]) log_p = log_p . view ( - 1 , num_samples , * log_p . size ()[ 1 :]) return z , log_q , log_p __init__ ( prior , q0 = distributions . Dirac (), flows = None , decoder = None ) Constructor of normalizing flow model Parameters: Name Type Description Default prior Prior distribution of te VAE, i.e. Gaussian required decoder Optional decoder None flows Flows to transform output of base encoder None q0 Base Encoder distributions.Dirac() Source code in normflows/core.py 457 458 459 460 461 462 463 464 465 466 467 468 469 470 def __init__ ( self , prior , q0 = distributions . Dirac (), flows = None , decoder = None ): \"\"\"Constructor of normalizing flow model Args: prior: Prior distribution of te VAE, i.e. Gaussian decoder: Optional decoder flows: Flows to transform output of base encoder q0: Base Encoder \"\"\" super () . __init__ () self . prior = prior self . decoder = decoder self . flows = nn . ModuleList ( flows ) self . q0 = q0 forward ( x , num_samples = 1 ) Takes data batch, samples num_samples for each data point from base distribution Parameters: Name Type Description Default x data batch required num_samples number of samples to draw for each data point 1 Returns: Type Description latent variables for each batch and sample, log_q, and log_p Source code in normflows/core.py 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 def forward ( self , x , num_samples = 1 ): \"\"\"Takes data batch, samples num_samples for each data point from base distribution Args: x: data batch num_samples: number of samples to draw for each data point Returns: latent variables for each batch and sample, log_q, and log_p \"\"\" z , log_q = self . q0 ( x , num_samples = num_samples ) # Flatten batch and sample dim z = z . view ( - 1 , * z . size ()[ 2 :]) log_q = log_q . view ( - 1 , * log_q . size ()[ 2 :]) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . prior . log_prob ( z ) if self . decoder is not None : log_p += self . decoder . log_prob ( x , z ) # Separate batch and sample dimension again z = z . view ( - 1 , num_samples , * z . size ()[ 1 :]) log_q = log_q . view ( - 1 , num_samples , * log_q . size ()[ 1 :]) log_p = log_p . view ( - 1 , num_samples , * log_p . size ()[ 1 :]) return z , log_q , log_p core_test distributions base AffineGaussian Bases: BaseDistribution Diagonal Gaussian an affine constant transformation applied to it, can be class conditional or not Source code in normflows/distributions/base.py 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 class AffineGaussian ( BaseDistribution ): \"\"\" Diagonal Gaussian an affine constant transformation applied to it, can be class conditional or not \"\"\" def __init__ ( self , shape , affine_shape , num_classes = None ): \"\"\"Constructor Args: shape: Shape of the variables affine_shape: Shape of the parameters in the affine transformation num_classes: Number of classes if the base is class conditional, None otherwise \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . affine_shape = affine_shape self . num_classes = num_classes self . class_cond = num_classes is not None # Affine transformation if self . class_cond : self . transform = flows . CCAffineConst ( self . affine_shape , self . num_classes ) else : self . transform = flows . AffineConstFlow ( self . affine_shape ) # Temperature parameter for annealed sampling self . temperature = None def forward ( self , num_samples = 1 , y = None ): dtype = self . transform . s . dtype device = self . transform . s . device if self . class_cond : if y is not None : num_samples = len ( y ) else : y = torch . randint ( self . num_classes , ( num_samples ,), device = device ) if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = dtype , device = device ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot if self . temperature is not None : log_scale = np . log ( self . temperature ) else : log_scale = 0.0 # Sample eps = torch . randn (( num_samples ,) + self . shape , dtype = dtype , device = device ) z = np . exp ( log_scale ) * eps # Get log prob log_p = ( - 0.5 * self . d * np . log ( 2 * np . pi ) - self . d * log_scale - 0.5 * torch . sum ( torch . pow ( eps , 2 ), dim = self . sum_dim ) ) # Apply transform if self . class_cond : z , log_det = self . transform ( z , y ) else : z , log_det = self . transform ( z ) log_p -= log_det return z , log_p def log_prob ( self , z , y = None ): # Perpare onehot encoding of class if needed if self . class_cond : if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = self . transform . s . dtype , device = self . transform . s . device , ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot if self . temperature is not None : log_scale = np . log ( self . temperature ) else : log_scale = 0.0 # Get log prob if self . class_cond : z , log_p = self . transform . inverse ( z , y ) else : z , log_p = self . transform . inverse ( z ) z = z / np . exp ( log_scale ) log_p = ( log_p - self . d * log_scale - 0.5 * self . d * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . pow ( z , 2 ), dim = self . sum_dim ) ) return log_p __init__ ( shape , affine_shape , num_classes = None ) Constructor Parameters: Name Type Description Default shape Shape of the variables required affine_shape Shape of the parameters in the affine transformation required num_classes Number of classes if the base is class conditional, None otherwise None Source code in normflows/distributions/base.py 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 def __init__ ( self , shape , affine_shape , num_classes = None ): \"\"\"Constructor Args: shape: Shape of the variables affine_shape: Shape of the parameters in the affine transformation num_classes: Number of classes if the base is class conditional, None otherwise \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . affine_shape = affine_shape self . num_classes = num_classes self . class_cond = num_classes is not None # Affine transformation if self . class_cond : self . transform = flows . CCAffineConst ( self . affine_shape , self . num_classes ) else : self . transform = flows . AffineConstFlow ( self . affine_shape ) # Temperature parameter for annealed sampling self . temperature = None BaseDistribution Bases: nn . Module Base distribution of a flow-based model Parameters do not depend of target variable (as is the case for a VAE encoder) Source code in normflows/distributions/base.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class BaseDistribution ( nn . Module ): \"\"\" Base distribution of a flow-based model Parameters do not depend of target variable (as is the case for a VAE encoder) \"\"\" def __init__ ( self ): super () . __init__ () def forward ( self , num_samples = 1 ): \"\"\"Samples from base distribution and calculates log probability Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution, log probability \"\"\" raise NotImplementedError def log_prob ( self , z ): \"\"\"Calculate log probability of batch of samples Args: z: Batch of random variables to determine log probability for Returns: log probability for each batch element \"\"\" raise NotImplementedError def sample ( self , num_samples = 1 , ** kwargs ): \"\"\"Samples from base distribution Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution \"\"\" z , _ = self . forward ( num_samples , ** kwargs ) return z forward ( num_samples = 1 ) Samples from base distribution and calculates log probability Parameters: Name Type Description Default num_samples Number of samples to draw from the distriubtion 1 Returns: Type Description Samples drawn from the distribution, log probability Source code in normflows/distributions/base.py 17 18 19 20 21 22 23 24 25 26 def forward ( self , num_samples = 1 ): \"\"\"Samples from base distribution and calculates log probability Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution, log probability \"\"\" raise NotImplementedError log_prob ( z ) Calculate log probability of batch of samples Parameters: Name Type Description Default z Batch of random variables to determine log probability for required Returns: Type Description log probability for each batch element Source code in normflows/distributions/base.py 28 29 30 31 32 33 34 35 36 37 def log_prob ( self , z ): \"\"\"Calculate log probability of batch of samples Args: z: Batch of random variables to determine log probability for Returns: log probability for each batch element \"\"\" raise NotImplementedError sample ( num_samples = 1 , ** kwargs ) Samples from base distribution Parameters: Name Type Description Default num_samples Number of samples to draw from the distriubtion 1 Returns: Type Description Samples drawn from the distribution Source code in normflows/distributions/base.py 39 40 41 42 43 44 45 46 47 48 49 def sample ( self , num_samples = 1 , ** kwargs ): \"\"\"Samples from base distribution Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution \"\"\" z , _ = self . forward ( num_samples , ** kwargs ) return z ClassCondDiagGaussian Bases: BaseDistribution Class conditional multivariate Gaussian distribution with diagonal covariance matrix Source code in normflows/distributions/base.py 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 class ClassCondDiagGaussian ( BaseDistribution ): \"\"\" Class conditional multivariate Gaussian distribution with diagonal covariance matrix \"\"\" def __init__ ( self , shape , num_classes ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension num_classes: Number of classes to condition on \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . perm = [ self . n_dim ] + list ( range ( self . n_dim )) self . d = np . prod ( shape ) self . num_classes = num_classes self . loc = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . log_scale = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . temperature = None # Temperature parameter for annealed sampling def forward ( self , num_samples = 1 , y = None ): if y is not None : num_samples = len ( y ) else : y = torch . randint ( self . num_classes , ( num_samples ,), device = self . loc . device ) if y . dim () == 1 : y_onehot = torch . zeros ( ( self . num_classes , num_samples ), dtype = self . loc . dtype , device = self . loc . device , ) y_onehot . scatter_ ( 0 , y [ None ], 1 ) y = y_onehot else : y = y . t () eps = torch . randn ( ( num_samples ,) + self . shape , dtype = self . loc . dtype , device = self . loc . device ) loc = ( self . loc @ y ) . permute ( * self . perm ) log_scale = ( self . log_scale @ y ) . permute ( * self . perm ) if self . temperature is not None : log_scale = np . log ( self . temperature ) + log_scale z = loc + torch . exp ( log_scale ) * eps log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow ( eps , 2 ), list ( range ( 1 , self . n_dim + 1 )) ) return z , log_p def log_prob ( self , z , y ): if y . dim () == 1 : y_onehot = torch . zeros ( ( self . num_classes , len ( y )), dtype = self . loc . dtype , device = self . loc . device ) y_onehot . scatter_ ( 0 , y [ None ], 1 ) y = y_onehot else : y = y . t () loc = ( self . loc @ y ) . permute ( * self . perm ) log_scale = ( self . log_scale @ y ) . permute ( * self . perm ) if self . temperature is not None : log_scale = np . log ( self . temperature ) + log_scale log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow (( z - loc ) / torch . exp ( log_scale ), 2 ), list ( range ( 1 , self . n_dim + 1 )), ) return log_p __init__ ( shape , num_classes ) Constructor Parameters: Name Type Description Default shape Tuple with shape of data, if int shape has one dimension required num_classes Number of classes to condition on required Source code in normflows/distributions/base.py 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 def __init__ ( self , shape , num_classes ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension num_classes: Number of classes to condition on \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . perm = [ self . n_dim ] + list ( range ( self . n_dim )) self . d = np . prod ( shape ) self . num_classes = num_classes self . loc = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . log_scale = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . temperature = None # Temperature parameter for annealed sampling DiagGaussian Bases: BaseDistribution Multivariate Gaussian distribution with diagonal covariance matrix Source code in normflows/distributions/base.py 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 class DiagGaussian ( BaseDistribution ): \"\"\" Multivariate Gaussian distribution with diagonal covariance matrix \"\"\" def __init__ ( self , shape , trainable = True ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) if trainable : self . loc = nn . Parameter ( torch . zeros ( 1 , * self . shape )) self . log_scale = nn . Parameter ( torch . zeros ( 1 , * self . shape )) else : self . register_buffer ( \"loc\" , torch . zeros ( 1 , * self . shape )) self . register_buffer ( \"log_scale\" , torch . zeros ( 1 , * self . shape )) self . temperature = None # Temperature parameter for annealed sampling def forward ( self , num_samples = 1 ): eps = torch . randn ( ( num_samples ,) + self . shape , dtype = self . loc . dtype , device = self . loc . device ) if self . temperature is None : log_scale = self . log_scale else : log_scale = self . log_scale + np . log ( self . temperature ) z = self . loc + torch . exp ( log_scale ) * eps log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow ( eps , 2 ), list ( range ( 1 , self . n_dim + 1 )) ) return z , log_p def log_prob ( self , z ): if self . temperature is None : log_scale = self . log_scale else : log_scale = self . log_scale + np . log ( self . temperature ) log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow (( z - self . loc ) / torch . exp ( log_scale ), 2 ), list ( range ( 1 , self . n_dim + 1 )), ) return log_p __init__ ( shape , trainable = True ) Constructor Parameters: Name Type Description Default shape Tuple with shape of data, if int shape has one dimension required Source code in normflows/distributions/base.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 def __init__ ( self , shape , trainable = True ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) if trainable : self . loc = nn . Parameter ( torch . zeros ( 1 , * self . shape )) self . log_scale = nn . Parameter ( torch . zeros ( 1 , * self . shape )) else : self . register_buffer ( \"loc\" , torch . zeros ( 1 , * self . shape )) self . register_buffer ( \"log_scale\" , torch . zeros ( 1 , * self . shape )) self . temperature = None # Temperature parameter for annealed sampling GaussianMixture Bases: BaseDistribution Mixture of Gaussians with diagonal covariance matrix Source code in normflows/distributions/base.py 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 class GaussianMixture ( BaseDistribution ): \"\"\" Mixture of Gaussians with diagonal covariance matrix \"\"\" def __init__ ( self , n_modes , dim , loc = None , scale = None , weights = None , trainable = True ): \"\"\"Constructor Args: n_modes: Number of modes of the mixture model dim: Number of dimensions of each Gaussian loc: List of mean values scale: List of diagonals of the covariance matrices weights: List of mode probabilities trainable: Flag, if true parameters will be optimized during training \"\"\" super () . __init__ () self . n_modes = n_modes self . dim = dim if loc is None : loc = np . random . randn ( self . n_modes , self . dim ) loc = np . array ( loc )[ None , ... ] if scale is None : scale = np . ones (( self . n_modes , self . dim )) scale = np . array ( scale )[ None , ... ] if weights is None : weights = np . ones ( self . n_modes ) weights = np . array ( weights )[ None , ... ] weights /= weights . sum ( 1 ) if trainable : self . loc = nn . Parameter ( torch . tensor ( 1.0 * loc )) self . log_scale = nn . Parameter ( torch . tensor ( np . log ( 1.0 * scale ))) self . weight_scores = nn . Parameter ( torch . tensor ( np . log ( 1.0 * weights ))) else : self . register_buffer ( \"loc\" , torch . tensor ( 1.0 * loc )) self . register_buffer ( \"log_scale\" , torch . tensor ( np . log ( 1.0 * scale ))) self . register_buffer ( \"weight_scores\" , torch . tensor ( np . log ( 1.0 * weights ))) def forward ( self , num_samples = 1 ): # Get weights weights = torch . softmax ( self . weight_scores , 1 ) # Sample mode indices mode = torch . multinomial ( weights [ 0 , :], num_samples , replacement = True ) mode_1h = nn . functional . one_hot ( mode , self . n_modes ) mode_1h = mode_1h [ ... , None ] # Get samples eps_ = torch . randn ( num_samples , self . dim , dtype = self . loc . dtype , device = self . loc . device ) scale_sample = torch . sum ( torch . exp ( self . log_scale ) * mode_1h , 1 ) loc_sample = torch . sum ( self . loc * mode_1h , 1 ) z = eps_ * scale_sample + loc_sample # Compute log probability eps = ( z [:, None , :] - self . loc ) / torch . exp ( self . log_scale ) log_p = ( - 0.5 * self . dim * np . log ( 2 * np . pi ) + torch . log ( weights ) - 0.5 * torch . sum ( torch . pow ( eps , 2 ), 2 ) - torch . sum ( self . log_scale , 2 ) ) log_p = torch . logsumexp ( log_p , 1 ) return z , log_p def log_prob ( self , z ): # Get weights weights = torch . softmax ( self . weight_scores , 1 ) # Compute log probability eps = ( z [:, None , :] - self . loc ) / torch . exp ( self . log_scale ) log_p = ( - 0.5 * self . dim * np . log ( 2 * np . pi ) + torch . log ( weights ) - 0.5 * torch . sum ( torch . pow ( eps , 2 ), 2 ) - torch . sum ( self . log_scale , 2 ) ) log_p = torch . logsumexp ( log_p , 1 ) return log_p __init__ ( n_modes , dim , loc = None , scale = None , weights = None , trainable = True ) Constructor Parameters: Name Type Description Default n_modes Number of modes of the mixture model required dim Number of dimensions of each Gaussian required loc List of mean values None scale List of diagonals of the covariance matrices None weights List of mode probabilities None trainable Flag, if true parameters will be optimized during training True Source code in normflows/distributions/base.py 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 def __init__ ( self , n_modes , dim , loc = None , scale = None , weights = None , trainable = True ): \"\"\"Constructor Args: n_modes: Number of modes of the mixture model dim: Number of dimensions of each Gaussian loc: List of mean values scale: List of diagonals of the covariance matrices weights: List of mode probabilities trainable: Flag, if true parameters will be optimized during training \"\"\" super () . __init__ () self . n_modes = n_modes self . dim = dim if loc is None : loc = np . random . randn ( self . n_modes , self . dim ) loc = np . array ( loc )[ None , ... ] if scale is None : scale = np . ones (( self . n_modes , self . dim )) scale = np . array ( scale )[ None , ... ] if weights is None : weights = np . ones ( self . n_modes ) weights = np . array ( weights )[ None , ... ] weights /= weights . sum ( 1 ) if trainable : self . loc = nn . Parameter ( torch . tensor ( 1.0 * loc )) self . log_scale = nn . Parameter ( torch . tensor ( np . log ( 1.0 * scale ))) self . weight_scores = nn . Parameter ( torch . tensor ( np . log ( 1.0 * weights ))) else : self . register_buffer ( \"loc\" , torch . tensor ( 1.0 * loc )) self . register_buffer ( \"log_scale\" , torch . tensor ( np . log ( 1.0 * scale ))) self . register_buffer ( \"weight_scores\" , torch . tensor ( np . log ( 1.0 * weights ))) GaussianPCA Bases: BaseDistribution Gaussian distribution resulting from linearly mapping a normal distributed latent variable describing the \"content of the target\" Source code in normflows/distributions/base.py 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 class GaussianPCA ( BaseDistribution ): \"\"\" Gaussian distribution resulting from linearly mapping a normal distributed latent variable describing the \"content of the target\" \"\"\" def __init__ ( self , dim , latent_dim = None , sigma = 0.1 ): \"\"\"Constructor Args: dim: Number of dimensions of the flow variables latent_dim: Number of dimensions of the latent \"content\" variable; if None it is set equal to dim sigma: Noise level \"\"\" super () . __init__ () self . dim = dim if latent_dim is None : self . latent_dim = dim else : self . latent_dim = latent_dim self . loc = nn . Parameter ( torch . zeros ( 1 , dim )) self . W = nn . Parameter ( torch . randn ( latent_dim , dim )) self . log_sigma = nn . Parameter ( torch . tensor ( np . log ( sigma ))) def forward ( self , num_samples = 1 ): eps = torch . randn ( num_samples , self . latent_dim , dtype = self . loc . dtype , device = self . loc . device ) z_ = torch . matmul ( eps , self . W ) z = z_ + self . loc Sig = torch . matmul ( self . W . T , self . W ) + torch . exp ( self . log_sigma * 2 ) * torch . eye ( self . dim , dtype = self . loc . dtype , device = self . loc . device ) log_p = ( self . dim / 2 * np . log ( 2 * np . pi ) - 0.5 * torch . det ( Sig ) - 0.5 * torch . sum ( z_ * torch . matmul ( z_ , torch . inverse ( Sig )), 1 ) ) return z , log_p def log_prob ( self , z ): z_ = z - self . loc Sig = torch . matmul ( self . W . T , self . W ) + torch . exp ( self . log_sigma * 2 ) * torch . eye ( self . dim , dtype = self . loc . dtype , device = self . loc . device ) log_p = ( self . dim / 2 * np . log ( 2 * np . pi ) - 0.5 * torch . det ( Sig ) - 0.5 * torch . sum ( z_ * torch . matmul ( z_ , torch . inverse ( Sig )), 1 ) ) return log_p __init__ ( dim , latent_dim = None , sigma = 0.1 ) Constructor Parameters: Name Type Description Default dim Number of dimensions of the flow variables required latent_dim Number of dimensions of the latent \"content\" variable; if None it is set equal to dim None sigma Noise level 0.1 Source code in normflows/distributions/base.py 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 def __init__ ( self , dim , latent_dim = None , sigma = 0.1 ): \"\"\"Constructor Args: dim: Number of dimensions of the flow variables latent_dim: Number of dimensions of the latent \"content\" variable; if None it is set equal to dim sigma: Noise level \"\"\" super () . __init__ () self . dim = dim if latent_dim is None : self . latent_dim = dim else : self . latent_dim = latent_dim self . loc = nn . Parameter ( torch . zeros ( 1 , dim )) self . W = nn . Parameter ( torch . randn ( latent_dim , dim )) self . log_sigma = nn . Parameter ( torch . tensor ( np . log ( sigma ))) GlowBase Bases: BaseDistribution Base distribution of the Glow model, i.e. Diagonal Gaussian with one mean and log scale for each channel Source code in normflows/distributions/base.py 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 class GlowBase ( BaseDistribution ): \"\"\" Base distribution of the Glow model, i.e. Diagonal Gaussian with one mean and log scale for each channel \"\"\" def __init__ ( self , shape , num_classes = None , logscale_factor = 3.0 ): \"\"\"Constructor Args: shape: Shape of the variables num_classes: Number of classes if the base is class conditional, None otherwise logscale_factor: Scaling factor for mean and log variance \"\"\" super () . __init__ () # Save shape and related statistics if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . num_pix = np . prod ( shape [ 1 :]) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . num_classes = num_classes self . class_cond = num_classes is not None self . logscale_factor = logscale_factor # Set up parameters self . loc = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . loc_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) # Class conditional parameter if needed if self . class_cond : self . loc_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ])) self . log_scale_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ]) ) # Temperature parameter for annealed sampling self . temperature = None def forward ( self , num_samples = 1 , y = None ): # Prepare parameter loc = self . loc * torch . exp ( self . loc_logs * self . logscale_factor ) log_scale = self . log_scale * torch . exp ( self . log_scale_logs * self . logscale_factor ) if self . class_cond : if y is not None : num_samples = len ( y ) else : y = torch . randint ( self . num_classes , ( num_samples ,), device = self . loc . device ) if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = self . loc . dtype , device = self . loc . device , ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot loc = loc + ( y @ self . loc_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) log_scale = log_scale + ( y @ self . log_scale_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) if self . temperature is not None : log_scale = log_scale + np . log ( self . temperature ) # Sample eps = torch . randn ( ( num_samples ,) + self . shape , dtype = self . loc . dtype , device = self . loc . device ) z = loc + torch . exp ( log_scale ) * eps # Get log prob log_p = ( - 0.5 * self . d * np . log ( 2 * np . pi ) - self . num_pix * torch . sum ( log_scale , dim = self . sum_dim ) - 0.5 * torch . sum ( torch . pow ( eps , 2 ), dim = self . sum_dim ) ) return z , log_p def log_prob ( self , z , y = None ): # Perpare parameter loc = self . loc * torch . exp ( self . loc_logs * self . logscale_factor ) log_scale = self . log_scale * torch . exp ( self . log_scale_logs * self . logscale_factor ) if self . class_cond : if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = self . loc . dtype , device = self . loc . device , ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot loc = loc + ( y @ self . loc_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) log_scale = log_scale + ( y @ self . log_scale_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) if self . temperature is not None : log_scale = log_scale + np . log ( self . temperature ) # Get log prob log_p = ( - 0.5 * self . d * np . log ( 2 * np . pi ) - self . num_pix * torch . sum ( log_scale , dim = self . sum_dim ) - 0.5 * torch . sum ( torch . pow (( z - loc ) / torch . exp ( log_scale ), 2 ), dim = self . sum_dim ) ) return log_p __init__ ( shape , num_classes = None , logscale_factor = 3.0 ) Constructor Parameters: Name Type Description Default shape Shape of the variables required num_classes Number of classes if the base is class conditional, None otherwise None logscale_factor Scaling factor for mean and log variance 3.0 Source code in normflows/distributions/base.py 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 def __init__ ( self , shape , num_classes = None , logscale_factor = 3.0 ): \"\"\"Constructor Args: shape: Shape of the variables num_classes: Number of classes if the base is class conditional, None otherwise logscale_factor: Scaling factor for mean and log variance \"\"\" super () . __init__ () # Save shape and related statistics if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . num_pix = np . prod ( shape [ 1 :]) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . num_classes = num_classes self . class_cond = num_classes is not None self . logscale_factor = logscale_factor # Set up parameters self . loc = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . loc_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) # Class conditional parameter if needed if self . class_cond : self . loc_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ])) self . log_scale_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ]) ) # Temperature parameter for annealed sampling self . temperature = None Uniform Bases: BaseDistribution Multivariate uniform distribution Source code in normflows/distributions/base.py 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 class Uniform ( BaseDistribution ): \"\"\" Multivariate uniform distribution \"\"\" def __init__ ( self , shape , low =- 1.0 , high = 1.0 ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension low: Lower bound of uniform distribution high: Upper bound of uniform distribution \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . d = np . prod ( shape ) self . low = torch . tensor ( low ) self . high = torch . tensor ( high ) self . log_prob_val = - self . d * np . log ( self . high - self . low ) def forward ( self , num_samples = 1 ): eps = torch . rand ( ( num_samples ,) + self . shape , dtype = self . low . dtype , device = self . low . device ) z = self . low + ( self . high - self . low ) * eps log_p = self . log_prob_val * torch . ones ( num_samples , device = self . low . device ) return z , log_p def log_prob ( self , z ): log_p = self . log_prob_val * torch . ones ( z . shape [ 0 ], device = z . device ) out_range = torch . logical_or ( z < self . low , z > self . high ) ind_inf = torch . any ( torch . reshape ( out_range , ( z . shape [ 0 ], - 1 )), dim =- 1 ) log_p [ ind_inf ] = - np . inf return log_p __init__ ( shape , low =- 1.0 , high = 1.0 ) Constructor Parameters: Name Type Description Default shape Tuple with shape of data, if int shape has one dimension required low Lower bound of uniform distribution -1.0 high Upper bound of uniform distribution 1.0 Source code in normflows/distributions/base.py 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 def __init__ ( self , shape , low =- 1.0 , high = 1.0 ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension low: Lower bound of uniform distribution high: Upper bound of uniform distribution \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . d = np . prod ( shape ) self . low = torch . tensor ( low ) self . high = torch . tensor ( high ) self . log_prob_val = - self . d * np . log ( self . high - self . low ) UniformGaussian Bases: BaseDistribution Distribution of a 1D random variable with some entries having a uniform and others a Gaussian distribution Source code in normflows/distributions/base.py 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 class UniformGaussian ( BaseDistribution ): \"\"\" Distribution of a 1D random variable with some entries having a uniform and others a Gaussian distribution \"\"\" def __init__ ( self , ndim , ind , scale = None ): \"\"\"Constructor Args: ndim: Int, number of dimensions ind: Iterable, indices of uniformly distributed entries scale: Iterable, standard deviation of Gaussian or width of uniform distribution \"\"\" super () . __init__ () self . ndim = ndim if isinstance ( ind , int ): ind = [ ind ] # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) if scale is None : self . register_buffer ( \"scale\" , torch . ones ( self . ndim )) else : self . register_buffer ( \"scale\" , scale ) def forward ( self , num_samples = 1 ): z = self . sample ( num_samples ) return z , self . log_prob ( z ) def sample ( self , num_samples = 1 ): eps_u = ( torch . rand ( ( num_samples , len ( self . ind )), dtype = self . scale . dtype , device = self . scale . device , ) - 0.5 ) eps_g = torch . randn ( ( num_samples , len ( self . ind_ )), dtype = self . scale . dtype , device = self . scale . device , ) z = torch . cat (( eps_u , eps_g ), - 1 ) z = z [ ... , self . inv_perm ] return self . scale * z def log_prob ( self , z ): log_p_u = torch . broadcast_to ( - torch . log ( self . scale [ self . ind ]), ( len ( z ), - 1 )) log_p_g = ( - 0.5 * np . log ( 2 * np . pi ) - torch . log ( self . scale [ self . ind_ ]) - 0.5 * torch . pow ( z [ ... , self . ind_ ] / self . scale [ self . ind_ ], 2 ) ) return torch . sum ( log_p_u , - 1 ) + torch . sum ( log_p_g , - 1 ) __init__ ( ndim , ind , scale = None ) Constructor Parameters: Name Type Description Default ndim Int, number of dimensions required ind Iterable, indices of uniformly distributed entries required scale Iterable, standard deviation of Gaussian or width of uniform distribution None Source code in normflows/distributions/base.py 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 def __init__ ( self , ndim , ind , scale = None ): \"\"\"Constructor Args: ndim: Int, number of dimensions ind: Iterable, indices of uniformly distributed entries scale: Iterable, standard deviation of Gaussian or width of uniform distribution \"\"\" super () . __init__ () self . ndim = ndim if isinstance ( ind , int ): ind = [ ind ] # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) if scale is None : self . register_buffer ( \"scale\" , torch . ones ( self . ndim )) else : self . register_buffer ( \"scale\" , scale ) base_test decoder BaseDecoder Bases: nn . Module Source code in normflows/distributions/decoder.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 class BaseDecoder ( nn . Module ): def __init__ ( self ): super () . __init__ () def forward ( self , z ): \"\"\"Decodes z to x Args: z: latent variable Returns: x, std of x \"\"\" raise NotImplementedError def log_prob ( self , x , z ): \"\"\"Log probability Args: x: observable z: latent variable Returns: log(p) of x given z \"\"\" raise NotImplementedError forward ( z ) Decodes z to x Parameters: Name Type Description Default z latent variable required Returns: Type Description x, std of x Source code in normflows/distributions/decoder.py 10 11 12 13 14 15 16 17 18 19 def forward ( self , z ): \"\"\"Decodes z to x Args: z: latent variable Returns: x, std of x \"\"\" raise NotImplementedError log_prob ( x , z ) Log probability Parameters: Name Type Description Default x observable required z latent variable required Returns: Type Description log(p) of x given z Source code in normflows/distributions/decoder.py 21 22 23 24 25 26 27 28 29 30 31 def log_prob ( self , x , z ): \"\"\"Log probability Args: x: observable z: latent variable Returns: log(p) of x given z \"\"\" raise NotImplementedError NNBernoulliDecoder Bases: BaseDecoder BaseDecoder representing a Bernoulli distribution with mean parametrized by a NN Source code in normflows/distributions/decoder.py 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 class NNBernoulliDecoder ( BaseDecoder ): \"\"\" BaseDecoder representing a Bernoulli distribution with mean parametrized by a NN \"\"\" def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) \"\"\" super () . __init__ () self . net = net def forward ( self , z ): mean = torch . sigmoid ( self . net ( z )) return mean def log_prob ( self , x , z ): score = self . net ( z ) if len ( z ) > len ( x ): x = x . unsqueeze ( 1 ) x = x . repeat ( 1 , z . size ()[ 0 ] // x . size ()[ 0 ], * (( x . dim () - 2 ) * [ 1 ])) . view ( - 1 , * x . size ()[ 2 :] ) log_sig = lambda a : - torch . relu ( - a ) - torch . log ( 1 + torch . exp ( - torch . abs ( a ))) log_p = torch . sum ( x * log_sig ( score ) + ( 1 - x ) * log_sig ( - score ), list ( range ( 1 , x . dim ())) ) return log_p __init__ ( net ) Constructor Parameters: Name Type Description Default net neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) required Source code in normflows/distributions/decoder.py 78 79 80 81 82 83 84 85 def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) \"\"\" super () . __init__ () self . net = net NNDiagGaussianDecoder Bases: BaseDecoder BaseDecoder representing a diagonal Gaussian distribution with mean and std parametrized by a NN Source code in normflows/distributions/decoder.py 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 class NNDiagGaussianDecoder ( BaseDecoder ): \"\"\" BaseDecoder representing a diagonal Gaussian distribution with mean and std parametrized by a NN \"\"\" def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean and standard deviation of diagonal Gaussian \"\"\" super () . __init__ () self . net = net def forward ( self , z ): mean_std = self . net ( z ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] std = torch . exp ( 0.5 * mean_std [:, n_hidden :, ... ]) return mean , std def log_prob ( self , x , z ): mean_std = self . net ( z ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] var = torch . exp ( mean_std [:, n_hidden :, ... ]) if len ( z ) > len ( x ): x = x . unsqueeze ( 1 ) x = x . repeat ( 1 , z . size ()[ 0 ] // x . size ()[ 0 ], * (( x . dim () - 2 ) * [ 1 ])) . view ( - 1 , * x . size ()[ 2 :] ) log_p = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 1 :])) * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . log ( var ) + ( x - mean ) ** 2 / var , list ( range ( 1 , z . dim ())) ) return log_p __init__ ( net ) Constructor Parameters: Name Type Description Default net neural network parametrizing mean and standard deviation of diagonal Gaussian required Source code in normflows/distributions/decoder.py 39 40 41 42 43 44 45 46 def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean and standard deviation of diagonal Gaussian \"\"\" super () . __init__ () self . net = net decoder_test distribution_test DistributionTest Bases: unittest . TestCase Generic test case for distribution modules Source code in normflows/distributions/distribution_test.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 class DistributionTest ( unittest . TestCase ): \"\"\" Generic test case for distribution modules \"\"\" def assertClose ( self , actual , expected , atol = None , rtol = None ): assert_close ( actual , expected , atol = atol , rtol = rtol ) def checkForward ( self , distribution , num_samples = 1 , ** kwargs ): # Do forward outputs , log_p = distribution ( num_samples , ** kwargs ) # Check type assert outputs . dtype == log_p . dtype # Check shape assert log_p . shape [ 0 ] == num_samples assert outputs . shape [ 0 ] == num_samples # Check dim assert outputs . dim () > log_p . dim () # Return results return outputs , log_p def checkLogProb ( self , distribution , inputs , ** kwargs ): # Compute log prob log_p = distribution . log_prob ( inputs , ** kwargs ) # Check type assert log_p . dtype == inputs . dtype # Check shape assert log_p . shape [ 0 ] == inputs . shape [ 0 ] # Return results return log_p def checkSample ( self , distribution , num_samples = 1 , ** kwargs ): # Do forward outputs = distribution . sample ( num_samples , ** kwargs ) # Check shape assert outputs . shape [ 0 ] == num_samples # Check dim assert outputs . dim () > 1 # Return results return outputs def checkForwardLogProb ( self , distribution , num_samples = 1 , atol = None , rtol = None , ** kwargs ): # Check forward outputs , log_p = self . checkForward ( distribution , num_samples , ** kwargs ) # Check log prob log_p_ = self . checkLogProb ( distribution , outputs , ** kwargs ) # Check consistency self . assertClose ( log_p_ , log_p , atol , rtol ) encoder BaseEncoder Bases: nn . Module Base distribution of a flow-based variational autoencoder Parameters of the distribution depend of the target variable x Source code in normflows/distributions/encoder.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 class BaseEncoder ( nn . Module ): \"\"\" Base distribution of a flow-based variational autoencoder Parameters of the distribution depend of the target variable x \"\"\" def __init__ ( self ): super () . __init__ () def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, first dimension is batch size num_samples: number of samples to draw per element of mini-batch Returns sample of z for x, log probability for sample \"\"\" raise NotImplementedError def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch size x: Variable to condition on, first dimension is batch size Returns: log probability of z given x \"\"\" raise NotImplementedError forward ( x , num_samples = 1 ) Parameters: Name Type Description Default x Variable to condition on, first dimension is batch size required num_samples number of samples to draw per element of mini-batch 1 Returns sample of z for x, log probability for sample Source code in normflows/distributions/encoder.py 15 16 17 18 19 20 21 22 23 24 def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, first dimension is batch size num_samples: number of samples to draw per element of mini-batch Returns sample of z for x, log probability for sample \"\"\" raise NotImplementedError log_prob ( z , x ) Parameters: Name Type Description Default z Primary random variable, first dimension is batch size required x Variable to condition on, first dimension is batch size required Returns: Type Description log probability of z given x Source code in normflows/distributions/encoder.py 26 27 28 29 30 31 32 33 34 35 36 def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch size x: Variable to condition on, first dimension is batch size Returns: log probability of z given x \"\"\" raise NotImplementedError ConstDiagGaussian Bases: BaseEncoder Source code in normflows/distributions/encoder.py 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 class ConstDiagGaussian ( BaseEncoder ): def __init__ ( self , loc , scale ): \"\"\"Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x Args: loc: mean vector of the distribution scale: vector of the standard deviations on the diagonal of the covariance matrix \"\"\" super () . __init__ () self . d = len ( loc ) if not torch . is_tensor ( loc ): loc = torch . tensor ( loc ) if not torch . is_tensor ( scale ): scale = torch . tensor ( scale ) self . loc = nn . Parameter ( loc . reshape (( 1 , 1 , self . d ))) self . scale = nn . Parameter ( scale ) def forward ( self , x = None , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, will only be used to determine the batch size num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" if x is not None : batch_size = len ( x ) else : batch_size = 1 eps = torch . randn (( batch_size , num_samples , self . d ), device = x . device ) z = self . loc + self . scale * eps log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * torch . pow ( eps , 2 ), 2 ) return z , log_q def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * (( z - self . loc ) / self . scale ) ** 2 , 2 ) return log_q __init__ ( loc , scale ) Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x Parameters: Name Type Description Default loc mean vector of the distribution required scale vector of the standard deviations on the diagonal of the covariance matrix required Source code in normflows/distributions/encoder.py 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 def __init__ ( self , loc , scale ): \"\"\"Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x Args: loc: mean vector of the distribution scale: vector of the standard deviations on the diagonal of the covariance matrix \"\"\" super () . __init__ () self . d = len ( loc ) if not torch . is_tensor ( loc ): loc = torch . tensor ( loc ) if not torch . is_tensor ( scale ): scale = torch . tensor ( scale ) self . loc = nn . Parameter ( loc . reshape (( 1 , 1 , self . d ))) self . scale = nn . Parameter ( scale ) forward ( x = None , num_samples = 1 ) Parameters: Name Type Description Default x Variable to condition on, will only be used to determine the batch size None num_samples number of samples to draw per element of mini-batch 1 Returns: Type Description sample of z for x, log probability for sample Source code in normflows/distributions/encoder.py 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 def forward ( self , x = None , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, will only be used to determine the batch size num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" if x is not None : batch_size = len ( x ) else : batch_size = 1 eps = torch . randn (( batch_size , num_samples , self . d ), device = x . device ) z = self . loc + self . scale * eps log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * torch . pow ( eps , 2 ), 2 ) return z , log_q log_prob ( z , x ) Parameters: Name Type Description Default z Primary random variable, first dimension is batch dimension required x Variable to condition on, first dimension is batch dimension required Returns: Type Description log probability of z given x Source code in normflows/distributions/encoder.py 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * (( z - self . loc ) / self . scale ) ** 2 , 2 ) return log_q NNDiagGaussian Bases: BaseEncoder Diagonal Gaussian distribution with mean and variance determined by a neural network Source code in normflows/distributions/encoder.py 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 class NNDiagGaussian ( BaseEncoder ): \"\"\" Diagonal Gaussian distribution with mean and variance determined by a neural network \"\"\" def __init__ ( self , net ): \"\"\"Construtor Args: net: net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) \"\"\" super () . __init__ () self . net = net def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" batch_size = len ( x ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) std = torch . exp ( 0.5 * mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) eps = torch . randn ( ( batch_size , num_samples ) + tuple ( mean . size ()[ 2 :]), device = x . device ) z = mean + std * eps log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( std ) + 0.5 * torch . pow ( eps , 2 ), list ( range ( 2 , z . dim ()))) return z , log_q def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) var = torch . exp ( mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . log ( var ) + ( z - mean ) ** 2 / var , 2 ) return log_q __init__ ( net ) Construtor Parameters: Name Type Description Default net net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) required Source code in normflows/distributions/encoder.py 135 136 137 138 139 140 141 142 def __init__ ( self , net ): \"\"\"Construtor Args: net: net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) \"\"\" super () . __init__ () self . net = net forward ( x , num_samples = 1 ) Parameters: Name Type Description Default x Variable to condition on required num_samples number of samples to draw per element of mini-batch 1 Returns: Type Description sample of z for x, log probability for sample Source code in normflows/distributions/encoder.py 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" batch_size = len ( x ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) std = torch . exp ( 0.5 * mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) eps = torch . randn ( ( batch_size , num_samples ) + tuple ( mean . size ()[ 2 :]), device = x . device ) z = mean + std * eps log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( std ) + 0.5 * torch . pow ( eps , 2 ), list ( range ( 2 , z . dim ()))) return z , log_q log_prob ( z , x ) Parameters: Name Type Description Default z Primary random variable, first dimension is batch dimension required x Variable to condition on, first dimension is batch dimension required Returns: Type Description log probability of z given x Source code in normflows/distributions/encoder.py 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) var = torch . exp ( mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . log ( var ) + ( z - mean ) ** 2 / var , 2 ) return log_q encoder_test linear_interpolation LinearInterpolation Linear interpolation of two distributions in the log space Source code in normflows/distributions/linear_interpolation.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 class LinearInterpolation : \"\"\" Linear interpolation of two distributions in the log space \"\"\" def __init__ ( self , dist1 , dist2 , alpha ): \"\"\"Constructor Interpolation parameter alpha: ``` log_p = alpha * log_p_1 + (1 - alpha) * log_p_2 ``` Args: dist1: First distribution dist2: Second distribution alpha: Interpolation parameter \"\"\" self . alpha = alpha self . dist1 = dist1 self . dist2 = dist2 def log_prob ( self , z ): return self . alpha * self . dist1 . log_prob ( z ) + ( 1 - self . alpha ) * self . dist2 . log_prob ( z ) __init__ ( dist1 , dist2 , alpha ) Constructor Interpolation parameter alpha: log_p = alpha * log_p_1 + (1 - alpha) * log_p_2 Parameters: Name Type Description Default dist1 First distribution required dist2 Second distribution required alpha Interpolation parameter required Source code in normflows/distributions/linear_interpolation.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 def __init__ ( self , dist1 , dist2 , alpha ): \"\"\"Constructor Interpolation parameter alpha: ``` log_p = alpha * log_p_1 + (1 - alpha) * log_p_2 ``` Args: dist1: First distribution dist2: Second distribution alpha: Interpolation parameter \"\"\" self . alpha = alpha self . dist1 = dist1 self . dist2 = dist2 mh_proposal DiagGaussianProposal Bases: MHProposal Diagonal Gaussian distribution with previous value as mean as a proposal for Metropolis Hastings algorithm Source code in normflows/distributions/mh_proposal.py 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 class DiagGaussianProposal ( MHProposal ): \"\"\" Diagonal Gaussian distribution with previous value as mean as a proposal for Metropolis Hastings algorithm \"\"\" def __init__ ( self , shape , scale ): \"\"\"Constructor Args: shape: Shape of variables to sample scale: Standard deviation of distribution \"\"\" super () . __init__ () self . shape = shape self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu . unsqueeze ( 0 )) def sample ( self , z ): num_samples = len ( z ) eps = torch . randn (( num_samples ,) + self . shape , dtype = z . dtype , device = z . device ) z_ = eps * self . scale + z return z_ def log_prob ( self , z_ , z ): log_p = - 0.5 * np . prod ( self . shape ) * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * torch . pow (( z_ - z ) / self . scale , 2 ), list ( range ( 1 , z . dim ())), ) return log_p def forward ( self , z ): num_samples = len ( z ) eps = torch . randn (( num_samples ,) + self . shape , dtype = z . dtype , device = z . device ) z_ = eps * self . scale + z log_p_diff = torch . zeros ( num_samples , dtype = z . dtype , device = z . device ) return z_ , log_p_diff __init__ ( shape , scale ) Constructor Parameters: Name Type Description Default shape Shape of variables to sample required scale Standard deviation of distribution required Source code in normflows/distributions/mh_proposal.py 53 54 55 56 57 58 59 60 61 62 63 def __init__ ( self , shape , scale ): \"\"\"Constructor Args: shape: Shape of variables to sample scale: Standard deviation of distribution \"\"\" super () . __init__ () self . shape = shape self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu . unsqueeze ( 0 )) MHProposal Bases: nn . Module Proposal distribution for the Metropolis Hastings algorithm Source code in normflows/distributions/mh_proposal.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 class MHProposal ( nn . Module ): \"\"\" Proposal distribution for the Metropolis Hastings algorithm \"\"\" def __init__ ( self ): super () . __init__ () def sample ( self , z ): \"\"\" Sample new value based on previous z \"\"\" raise NotImplementedError def log_prob ( self , z_ , z ): \"\"\" Args: z_: Potential new sample z: Previous sample Returns: Log probability of proposal distribution \"\"\" raise NotImplementedError def forward ( self , z ): \"\"\"Draw samples given z and compute log probability difference ``` log(p(z | z_new)) - log(p(z_new | z)) ``` Args: z: Previous samples Returns: Proposal, difference of log probability ratio \"\"\" raise NotImplementedError forward ( z ) Draw samples given z and compute log probability difference log(p(z | z_new)) - log(p(z_new | z)) Parameters: Name Type Description Default z Previous samples required Returns: Type Description Proposal, difference of log probability ratio Source code in normflows/distributions/mh_proposal.py 31 32 33 34 35 36 37 38 39 40 41 42 43 44 def forward ( self , z ): \"\"\"Draw samples given z and compute log probability difference ``` log(p(z | z_new)) - log(p(z_new | z)) ``` Args: z: Previous samples Returns: Proposal, difference of log probability ratio \"\"\" raise NotImplementedError log_prob ( z_ , z ) Parameters: Name Type Description Default z_ Potential new sample required z Previous sample required Returns: Type Description Log probability of proposal distribution Source code in normflows/distributions/mh_proposal.py 20 21 22 23 24 25 26 27 28 29 def log_prob ( self , z_ , z ): \"\"\" Args: z_: Potential new sample z: Previous sample Returns: Log probability of proposal distribution \"\"\" raise NotImplementedError sample ( z ) Sample new value based on previous z Source code in normflows/distributions/mh_proposal.py 14 15 16 17 18 def sample ( self , z ): \"\"\" Sample new value based on previous z \"\"\" raise NotImplementedError prior ImagePrior Bases: nn . Module Intensities of an image determine probability density of prior Source code in normflows/distributions/prior.py 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 class ImagePrior ( nn . Module ): \"\"\" Intensities of an image determine probability density of prior \"\"\" def __init__ ( self , image , x_range = [ - 3 , 3 ], y_range = [ - 3 , 3 ], eps = 1.0e-10 ): \"\"\"Constructor Args: image: image as np matrix x_range: x range to position image at y_range: y range to position image at eps: small value to add to image to avoid log(0) problems \"\"\" super () . __init__ () image_ = np . flip ( image , 0 ) . transpose () + eps self . image_cpu = torch . tensor ( image_ / np . max ( image_ )) self . image_size_cpu = self . image_cpu . size () self . x_range = torch . tensor ( x_range ) self . y_range = torch . tensor ( y_range ) self . register_buffer ( \"image\" , self . image_cpu ) self . register_buffer ( \"image_size\" , torch . tensor ( self . image_size_cpu ) . unsqueeze ( 0 ) ) self . register_buffer ( \"density\" , torch . log ( self . image_cpu / torch . sum ( self . image_cpu )) ) self . register_buffer ( \"scale\" , torch . tensor ( [[ self . x_range [ 1 ] - self . x_range [ 0 ], self . y_range [ 1 ] - self . y_range [ 0 ]]] ), ) self . register_buffer ( \"shift\" , torch . tensor ([[ self . x_range [ 0 ], self . y_range [ 0 ]]]) ) def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" z_ = torch . clamp (( z - self . shift ) / self . scale , max = 1 , min = 0 ) ind = ( z_ * ( self . image_size - 1 )) . long () return self . density [ ind [:, 0 ], ind [:, 1 ]] def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" z_ = torch . rand ( ( num_steps , 2 ), dtype = self . image . dtype , device = self . image . device ) prob = torch . rand ( num_steps , dtype = self . image . dtype , device = self . image . device ) ind = ( z_ * ( self . image_size - 1 )) . long () intensity = self . image [ ind [:, 0 ], ind [:, 1 ]] accept = intensity > prob z = z_ [ accept , :] * self . scale + self . shift return z def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . ones (( 0 , 2 ), dtype = self . image . dtype , device = self . image . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z __init__ ( image , x_range = [ - 3 , 3 ], y_range = [ - 3 , 3 ], eps = 1e-10 ) Constructor Parameters: Name Type Description Default image image as np matrix required x_range x range to position image at [-3, 3] y_range y range to position image at [-3, 3] eps small value to add to image to avoid log(0) problems 1e-10 Source code in normflows/distributions/prior.py 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 def __init__ ( self , image , x_range = [ - 3 , 3 ], y_range = [ - 3 , 3 ], eps = 1.0e-10 ): \"\"\"Constructor Args: image: image as np matrix x_range: x range to position image at y_range: y range to position image at eps: small value to add to image to avoid log(0) problems \"\"\" super () . __init__ () image_ = np . flip ( image , 0 ) . transpose () + eps self . image_cpu = torch . tensor ( image_ / np . max ( image_ )) self . image_size_cpu = self . image_cpu . size () self . x_range = torch . tensor ( x_range ) self . y_range = torch . tensor ( y_range ) self . register_buffer ( \"image\" , self . image_cpu ) self . register_buffer ( \"image_size\" , torch . tensor ( self . image_size_cpu ) . unsqueeze ( 0 ) ) self . register_buffer ( \"density\" , torch . log ( self . image_cpu / torch . sum ( self . image_cpu )) ) self . register_buffer ( \"scale\" , torch . tensor ( [[ self . x_range [ 1 ] - self . x_range [ 0 ], self . y_range [ 1 ] - self . y_range [ 0 ]]] ), ) self . register_buffer ( \"shift\" , torch . tensor ([[ self . x_range [ 0 ], self . y_range [ 0 ]]]) ) log_prob ( z ) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 59 60 61 62 63 64 65 66 67 68 69 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" z_ = torch . clamp (( z - self . shift ) / self . scale , max = 1 , min = 0 ) ind = ( z_ * ( self . image_size - 1 )) . long () return self . density [ ind [:, 0 ], ind [:, 1 ]] rejection_sampling ( num_steps = 1 ) Perform rejection sampling on image distribution Parameters: Name Type Description Default num_steps Number of rejection sampling steps to perform 1 Returns: Type Description Accepted samples Source code in normflows/distributions/prior.py 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" z_ = torch . rand ( ( num_steps , 2 ), dtype = self . image . dtype , device = self . image . device ) prob = torch . rand ( num_steps , dtype = self . image . dtype , device = self . image . device ) ind = ( z_ * ( self . image_size - 1 )) . long () intensity = self . image [ ind [:, 0 ], ind [:, 1 ]] accept = intensity > prob z = z_ [ accept , :] * self . scale + self . shift return z sample ( num_samples = 1 ) Sample from image distribution through rejection sampling Parameters: Name Type Description Default num_samples Number of samples to draw 1 Returns: Type Description Samples Source code in normflows/distributions/prior.py 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . ones (( 0 , 2 ), dtype = self . image . dtype , device = self . image . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z PriorDistribution Source code in normflows/distributions/prior.py 6 7 8 9 10 11 12 13 14 15 16 17 18 class PriorDistribution : def __init__ ( self ): raise NotImplementedError def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError log_prob ( z ) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 10 11 12 13 14 15 16 17 18 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError Sinusoidal Bases: PriorDistribution Source code in normflows/distributions/prior.py 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 class Sinusoidal ( PriorDistribution ): def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density given by ``` w_1(z) = sin(2*pi / period * z[0]) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 ``` Args: scale: scale of the distribution, see formula period: period of the sinosoidal \"\"\" self . scale = scale self . period = period def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 w_1(z) = sin(2*pi / period * z[0]) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) log_prob = ( - 0.5 * (( z_ [ 1 ] - w_1 ( z_ )) / ( self . scale )) ** 2 - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) # add Gaussian envelope for valid p(z) return log_prob __init__ ( scale , period ) Distribution 2d with sinusoidal density given by w_1(z) = sin(2*pi / period * z[0]) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 Parameters: Name Type Description Default scale scale of the distribution, see formula required period period of the sinosoidal required Source code in normflows/distributions/prior.py 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density given by ``` w_1(z) = sin(2*pi / period * z[0]) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 ``` Args: scale: scale of the distribution, see formula period: period of the sinosoidal \"\"\" self . scale = scale self . period = period log_prob ( z ) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 w_1(z) = sin(2*pi / period * z[0]) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 w_1(z) = sin(2*pi / period * z[0]) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) log_prob = ( - 0.5 * (( z_ [ 1 ] - w_1 ( z_ )) / ( self . scale )) ** 2 - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) # add Gaussian envelope for valid p(z) return log_prob Sinusoidal_gap Bases: PriorDistribution Source code in normflows/distributions/prior.py 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 class Sinusoidal_gap ( PriorDistribution ): def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with gap given by ``` w_1(z) = sin(2*pi / period * z[0]) w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w2_scale = 0.6 self . w2_amp = 3.0 self . w2_mu = 1.0 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_2 = lambda x : self . w2_amp * torch . exp ( - 0.5 * (( z_ [ 0 ] - self . w2_mu ) / self . w2_scale ) ** 2 ) eps = torch . abs ( w_2 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_2 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / self . scale ) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob __init__ ( scale , period ) Distribution 2d with sinusoidal density with gap given by w_1(z) = sin(2*pi / period * z[0]) w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2)) Parameters: Name Type Description Default loc distance of modes from the origin required scale scale of modes required Source code in normflows/distributions/prior.py 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with gap given by ``` w_1(z) = sin(2*pi / period * z[0]) w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w2_scale = 0.6 self . w2_amp = 3.0 self . w2_mu = 1.0 log_prob ( z ) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_2 = lambda x : self . w2_amp * torch . exp ( - 0.5 * (( z_ [ 0 ] - self . w2_mu ) / self . w2_scale ) ** 2 ) eps = torch . abs ( w_2 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_2 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / self . scale ) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob Sinusoidal_split Bases: PriorDistribution Source code in normflows/distributions/prior.py 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 class Sinusoidal_split ( PriorDistribution ): def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with split given by ``` w_1(z) = sin(2*pi / period * z[0]) w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w3_scale = 0.3 self . w3_amp = 3.0 self . w3_mu = 1.0 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_3 = lambda x : self . w3_amp * torch . sigmoid ( ( z_ [ 0 ] - self . w3_mu ) / self . w3_scale ) eps = torch . abs ( w_3 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_3 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / ( self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob __init__ ( scale , period ) Distribution 2d with sinusoidal density with split given by w_1(z) = sin(2*pi / period * z[0]) w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2)) Parameters: Name Type Description Default loc distance of modes from the origin required scale scale of modes required Source code in normflows/distributions/prior.py 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with split given by ``` w_1(z) = sin(2*pi / period * z[0]) w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w3_scale = 0.3 self . w3_amp = 3.0 self . w3_mu = 1.0 log_prob ( z ) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_3 = lambda x : self . w3_amp * torch . sigmoid ( ( z_ [ 0 ] - self . w3_mu ) / self . w3_scale ) eps = torch . abs ( w_3 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_3 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / ( self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob Smiley Bases: PriorDistribution Source code in normflows/distributions/prior.py 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 class Smiley ( PriorDistribution ): def __init__ ( self , scale ): \"\"\"Distribution 2d of a smiley :) Args: scale: scale of the smiley \"\"\" self . scale = scale self . loc = 2.0 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z log_prob = ( - 0.5 * (( torch . norm ( z_ , dim = 0 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( torch . abs ( z_ [ 1 ] + 0.8 ) - 1.2 ) / ( 2 * self . scale )) ** 2 ) return log_prob __init__ ( scale ) Distribution 2d of a smiley :) Parameters: Name Type Description Default scale scale of the smiley required Source code in normflows/distributions/prior.py 300 301 302 303 304 305 306 307 def __init__ ( self , scale ): \"\"\"Distribution 2d of a smiley :) Args: scale: scale of the smiley \"\"\" self . scale = scale self . loc = 2.0 log_prob ( z ) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z log_prob = ( - 0.5 * (( torch . norm ( z_ , dim = 0 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( torch . abs ( z_ [ 1 ] + 0.8 ) - 1.2 ) / ( 2 * self . scale )) ** 2 ) return log_prob TwoModes Bases: PriorDistribution Source code in normflows/distributions/prior.py 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 class TwoModes ( PriorDistribution ): def __init__ ( self , loc , scale ): \"\"\"Distribution 2d with two modes Distribution 2d with two modes at ```z[0] = -loc``` and ```z[0] = loc``` following the density ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . loc = loc self . scale = scale def log_prob ( self , z ): \"\"\" ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) eps = torch . abs ( torch . tensor ( self . loc )) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( a - eps ) / ( 3 * self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( a * eps ) / ( 3 * self . scale ) ** 2 )) ) return log_prob __init__ ( loc , scale ) Distribution 2d with two modes Distribution 2d with two modes at z[0] = -loc and z[0] = loc following the density log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) Args: loc: distance of modes from the origin scale: scale of modes Source code in normflows/distributions/prior.py 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 def __init__ ( self , loc , scale ): \"\"\"Distribution 2d with two modes Distribution 2d with two modes at ```z[0] = -loc``` and ```z[0] = loc``` following the density ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . loc = loc self . scale = scale log_prob ( z ) log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 def log_prob ( self , z ): \"\"\" ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) eps = torch . abs ( torch . tensor ( self . loc )) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( a - eps ) / ( 3 * self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( a * eps ) / ( 3 * self . scale ) ** 2 )) ) return log_prob prior_test target CircularGaussianMixture Bases: nn . Module Two-dimensional Gaussian mixture arranged in a circle Source code in normflows/distributions/target.py 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 class CircularGaussianMixture ( nn . Module ): \"\"\" Two-dimensional Gaussian mixture arranged in a circle \"\"\" def __init__ ( self , n_modes = 8 ): \"\"\"Constructor Args: n_modes: Number of modes \"\"\" super ( CircularGaussianMixture , self ) . __init__ () self . n_modes = n_modes self . register_buffer ( \"scale\" , torch . tensor ( 2 / 3 * np . sin ( np . pi / self . n_modes )) . float () ) def log_prob ( self , z ): d = torch . zeros (( len ( z ), 0 ), dtype = z . dtype , device = z . device ) for i in range ( self . n_modes ): d_ = ( ( z [:, 0 ] - 2 * np . sin ( 2 * np . pi / self . n_modes * i )) ** 2 + ( z [:, 1 ] - 2 * np . cos ( 2 * np . pi / self . n_modes * i )) ** 2 ) / ( 2 * self . scale ** 2 ) d = torch . cat (( d , d_ [:, None ]), 1 ) log_p = - torch . log ( 2 * np . pi * self . scale ** 2 * self . n_modes ) + torch . logsumexp ( - d , 1 ) return log_p def sample ( self , num_samples = 1 ): eps = torch . randn ( ( num_samples , 2 ), dtype = self . scale . dtype , device = self . scale . device ) phi = ( 2 * np . pi / self . n_modes * torch . randint ( 0 , self . n_modes , ( num_samples ,), device = self . scale . device ) ) loc = torch . stack (( 2 * torch . sin ( phi ), 2 * torch . cos ( phi )), 1 ) . type ( eps . dtype ) return eps * self . scale + loc __init__ ( n_modes = 8 ) Constructor Parameters: Name Type Description Default n_modes Number of modes 8 Source code in normflows/distributions/target.py 137 138 139 140 141 142 143 144 145 146 147 def __init__ ( self , n_modes = 8 ): \"\"\"Constructor Args: n_modes: Number of modes \"\"\" super ( CircularGaussianMixture , self ) . __init__ () self . n_modes = n_modes self . register_buffer ( \"scale\" , torch . tensor ( 2 / 3 * np . sin ( np . pi / self . n_modes )) . float () ) RingMixture Bases: Target Mixture of ring distributions in two dimensions Source code in normflows/distributions/target.py 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 class RingMixture ( Target ): \"\"\" Mixture of ring distributions in two dimensions \"\"\" def __init__ ( self , n_rings = 2 ): super () . __init__ () self . n_dims = 2 self . max_log_prob = 0.0 self . n_rings = n_rings self . scale = 1 / 4 / self . n_rings def log_prob ( self , z ): d = torch . zeros (( len ( z ), 0 ), dtype = z . dtype , device = z . device ) for i in range ( self . n_rings ): d_ = (( torch . norm ( z , dim = 1 ) - 2 / self . n_rings * ( i + 1 )) ** 2 ) / ( 2 * self . scale ** 2 ) d = torch . cat (( d , d_ [:, None ]), 1 ) return torch . logsumexp ( - d , 1 ) Target Bases: nn . Module Sample target distributions to test models Source code in normflows/distributions/target.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 class Target ( nn . Module ): \"\"\" Sample target distributions to test models \"\"\" def __init__ ( self , prop_scale = torch . tensor ( 6.0 ), prop_shift = torch . tensor ( - 3.0 )): \"\"\"Constructor Args: prop_scale: Scale for the uniform proposal prop_shift: Shift for the uniform proposal \"\"\" super () . __init__ () self . register_buffer ( \"prop_scale\" , prop_scale ) self . register_buffer ( \"prop_shift\" , prop_shift ) def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError ( \"The log probability is not implemented yet.\" ) def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" eps = torch . rand ( ( num_steps , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device , ) z_ = self . prop_scale * eps + self . prop_shift prob = torch . rand ( num_steps , dtype = self . prop_scale . dtype , device = self . prop_scale . device ) prob_ = torch . exp ( self . log_prob ( z_ ) - self . max_log_prob ) accept = prob_ > prob z = z_ [ accept , :] return z def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . zeros ( ( 0 , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z __init__ ( prop_scale = torch . tensor ( 6.0 ), prop_shift = torch . tensor ( - 3.0 )) Constructor Parameters: Name Type Description Default prop_scale Scale for the uniform proposal torch.tensor(6.0) prop_shift Shift for the uniform proposal torch.tensor(-3.0) Source code in normflows/distributions/target.py 13 14 15 16 17 18 19 20 21 22 def __init__ ( self , prop_scale = torch . tensor ( 6.0 ), prop_shift = torch . tensor ( - 3.0 )): \"\"\"Constructor Args: prop_scale: Scale for the uniform proposal prop_shift: Shift for the uniform proposal \"\"\" super () . __init__ () self . register_buffer ( \"prop_scale\" , prop_scale ) self . register_buffer ( \"prop_shift\" , prop_shift ) log_prob ( z ) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/target.py 24 25 26 27 28 29 30 31 32 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError ( \"The log probability is not implemented yet.\" ) rejection_sampling ( num_steps = 1 ) Perform rejection sampling on image distribution Parameters: Name Type Description Default num_steps Number of rejection sampling steps to perform 1 Returns: Type Description Accepted samples Source code in normflows/distributions/target.py 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" eps = torch . rand ( ( num_steps , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device , ) z_ = self . prop_scale * eps + self . prop_shift prob = torch . rand ( num_steps , dtype = self . prop_scale . dtype , device = self . prop_scale . device ) prob_ = torch . exp ( self . log_prob ( z_ ) - self . max_log_prob ) accept = prob_ > prob z = z_ [ accept , :] return z sample ( num_samples = 1 ) Sample from image distribution through rejection sampling Parameters: Name Type Description Default num_samples Number of samples to draw 1 Returns: Type Description Samples Source code in normflows/distributions/target.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . zeros ( ( 0 , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z TwoIndependent Bases: Target Target distribution that combines two independent distributions of equal size into one distribution. This is needed for Augmented Normalizing Flows, see https://arxiv.org/abs/2002.07101 Source code in normflows/distributions/target.py 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 class TwoIndependent ( Target ): \"\"\" Target distribution that combines two independent distributions of equal size into one distribution. This is needed for Augmented Normalizing Flows, see https://arxiv.org/abs/2002.07101 \"\"\" def __init__ ( self , target1 , target2 ): super () . __init__ () self . target1 = target1 self . target2 = target2 self . split = Split ( mode = 'channel' ) def log_prob ( self , z ): z1 , z2 = self . split ( z )[ 0 ] return self . target1 . log_prob ( z1 ) + self . target2 . log_prob ( z2 ) def sample ( self , num_samples = 1 ): z1 = self . target1 . sample ( num_samples ) z2 = self . target2 . sample ( num_samples ) return self . split . inverse ([ z1 , z2 ])[ 0 ] TwoMoons Bases: Target Bimodal two-dimensional distribution Source code in normflows/distributions/target.py 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 class TwoMoons ( Target ): \"\"\" Bimodal two-dimensional distribution \"\"\" def __init__ ( self ): super () . __init__ () self . n_dims = 2 self . max_log_prob = 0.0 def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2 + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2) + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - 2 ) / 0.2 ) ** 2 - 0.5 * (( a - 2 ) / 0.3 ) ** 2 + torch . log ( 1 + torch . exp ( - 4 * a / 0.09 )) ) return log_prob log_prob ( z ) log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2 + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2) + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2)) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/target.py 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2 + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2) + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - 2 ) / 0.2 ) ** 2 - 0.5 * (( a - 2 ) / 0.3 ) ** 2 + torch . log ( 1 + torch . exp ( - 4 * a / 0.09 )) ) return log_prob target_test flows affine autoregressive Autoregressive Bases: Flow Transforms each input variable with an invertible elementwise transformation. The parameters of each invertible elementwise transformation can be functions of previous input variables, but they must not depend on the current or any following input variables. NOTE Calculating the inverse transform is D times slower than calculating the forward transform, where D is the dimensionality of the input to the transform. Source code in normflows/flows/affine/autoregressive.py 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class Autoregressive ( Flow ): \"\"\"Transforms each input variable with an invertible elementwise transformation. The parameters of each invertible elementwise transformation can be functions of previous input variables, but they must not depend on the current or any following input variables. **NOTE** Calculating the inverse transform is D times slower than calculating the forward transform, where D is the dimensionality of the input to the transform. \"\"\" def __init__ ( self , autoregressive_net ): super ( Autoregressive , self ) . __init__ () self . autoregressive_net = autoregressive_net def forward ( self , inputs , context = None ): autoregressive_params = self . autoregressive_net ( inputs , context ) outputs , logabsdet = self . _elementwise_forward ( inputs , autoregressive_params ) return outputs , logabsdet def inverse ( self , inputs , context = None ): num_inputs = np . prod ( inputs . shape [ 1 :]) outputs = torch . zeros_like ( inputs ) logabsdet = None for _ in range ( num_inputs ): autoregressive_params = self . autoregressive_net ( outputs , context ) outputs , logabsdet = self . _elementwise_inverse ( inputs , autoregressive_params ) return outputs , logabsdet def _output_dim_multiplier ( self ): raise NotImplementedError () def _elementwise_forward ( self , inputs , autoregressive_params ): raise NotImplementedError () def _elementwise_inverse ( self , inputs , autoregressive_params ): raise NotImplementedError () autoregressive_test coupling AffineConstFlow Bases: Flow scales and shifts with learned constants per dimension. In the NICE paper there is a scaling layer which is a special case of this where t is None Source code in normflows/flows/affine/coupling.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 class AffineConstFlow ( Flow ): \"\"\" scales and shifts with learned constants per dimension. In the NICE paper there is a scaling layer which is a special case of this where t is None \"\"\" def __init__ ( self , shape , scale = True , shift = True ): \"\"\"Constructor Args: shape: Shape of the coupling layer scale: Flag whether to apply scaling shift: Flag whether to apply shift logscale_factor: Optional factor which can be used to control the scale of the log scale factor \"\"\" super () . __init__ () if scale : self . s = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"s\" , torch . zeros ( shape )[ None ]) if shift : self . t = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"t\" , torch . zeros ( shape )[ None ]) self . n_dim = self . s . dim () self . batch_dims = torch . nonzero ( torch . tensor ( self . s . shape ) == 1 , as_tuple = False )[:, 0 ] . tolist () def forward ( self , z ): z_ = z * torch . exp ( self . s ) + self . t if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = prod_batch_dims * torch . sum ( self . s ) return z_ , log_det def inverse ( self , z ): z_ = ( z - self . t ) * torch . exp ( - self . s ) if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = - prod_batch_dims * torch . sum ( self . s ) return z_ , log_det __init__ ( shape , scale = True , shift = True ) Constructor Parameters: Name Type Description Default shape Shape of the coupling layer required scale Flag whether to apply scaling True shift Flag whether to apply shift True logscale_factor Optional factor which can be used to control the scale of the log scale factor required Source code in normflows/flows/affine/coupling.py 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 def __init__ ( self , shape , scale = True , shift = True ): \"\"\"Constructor Args: shape: Shape of the coupling layer scale: Flag whether to apply scaling shift: Flag whether to apply shift logscale_factor: Optional factor which can be used to control the scale of the log scale factor \"\"\" super () . __init__ () if scale : self . s = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"s\" , torch . zeros ( shape )[ None ]) if shift : self . t = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"t\" , torch . zeros ( shape )[ None ]) self . n_dim = self . s . dim () self . batch_dims = torch . nonzero ( torch . tensor ( self . s . shape ) == 1 , as_tuple = False )[:, 0 ] . tolist () AffineCoupling Bases: Flow Affine Coupling layer as introduced RealNVP paper, see arXiv: 1605.08803 Source code in normflows/flows/affine/coupling.py 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 class AffineCoupling ( Flow ): \"\"\" Affine Coupling layer as introduced RealNVP paper, see arXiv: 1605.08803 \"\"\" def __init__ ( self , param_map , scale = True , scale_map = \"exp\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model \"\"\" super () . __init__ () self . add_module ( \"param_map\" , param_map ) self . scale = scale self . scale_map = scale_map def forward ( self , z ): \"\"\" z is a list of z1 and z2; ```z = [z1, z2]``` z1 is left constant and affine map is applied to z2 with parameters depending on z1 Args: z \"\"\" z1 , z2 = z param = self . param_map ( z1 ) if self . scale : shift = param [:, 0 :: 2 , ... ] scale_ = param [:, 1 :: 2 , ... ] if self . scale_map == \"exp\" : z2 = z2 * torch . exp ( scale_ ) + shift log_det = torch . sum ( scale_ , dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 / scale + shift log_det = - torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid_inv\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 * scale + shift log_det = torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) else : raise NotImplementedError ( \"This scale map is not implemented.\" ) else : z2 = z2 + param log_det = zero_log_det_like_z ( z2 ) return [ z1 , z2 ], log_det def inverse ( self , z ): z1 , z2 = z param = self . param_map ( z1 ) if self . scale : shift = param [:, 0 :: 2 , ... ] scale_ = param [:, 1 :: 2 , ... ] if self . scale_map == \"exp\" : z2 = ( z2 - shift ) * torch . exp ( - scale_ ) log_det = - torch . sum ( scale_ , dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = ( z2 - shift ) * scale log_det = torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid_inv\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = ( z2 - shift ) / scale log_det = - torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) else : raise NotImplementedError ( \"This scale map is not implemented.\" ) else : z2 = z2 - param log_det = zero_log_det_like_z ( z2 ) return [ z1 , z2 ], log_det __init__ ( param_map , scale = True , scale_map = 'exp' ) Constructor Parameters: Name Type Description Default param_map Maps features to shift and scale parameter (if applicable) required scale Flag whether scale shall be applied True scale_map Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model 'exp' Source code in normflows/flows/affine/coupling.py 104 105 106 107 108 109 110 111 112 113 114 115 def __init__ ( self , param_map , scale = True , scale_map = \"exp\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model \"\"\" super () . __init__ () self . add_module ( \"param_map\" , param_map ) self . scale = scale self . scale_map = scale_map forward ( z ) z is a list of z1 and z2; z = [z1, z2] z1 is left constant and affine map is applied to z2 with parameters depending on z1 Source code in normflows/flows/affine/coupling.py 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 def forward ( self , z ): \"\"\" z is a list of z1 and z2; ```z = [z1, z2]``` z1 is left constant and affine map is applied to z2 with parameters depending on z1 Args: z \"\"\" z1 , z2 = z param = self . param_map ( z1 ) if self . scale : shift = param [:, 0 :: 2 , ... ] scale_ = param [:, 1 :: 2 , ... ] if self . scale_map == \"exp\" : z2 = z2 * torch . exp ( scale_ ) + shift log_det = torch . sum ( scale_ , dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 / scale + shift log_det = - torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid_inv\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 * scale + shift log_det = torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) else : raise NotImplementedError ( \"This scale map is not implemented.\" ) else : z2 = z2 + param log_det = zero_log_det_like_z ( z2 ) return [ z1 , z2 ], log_det AffineCouplingBlock Bases: Flow Affine Coupling layer including split and merge operation Source code in normflows/flows/affine/coupling.py 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 class AffineCouplingBlock ( Flow ): \"\"\" Affine Coupling layer including split and merge operation \"\"\" def __init__ ( self , param_map , scale = True , scale_map = \"exp\" , split_mode = \"channel\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Split layer self . flows += [ Split ( split_mode )] # Affine coupling layer self . flows += [ AffineCoupling ( param_map , scale , scale_map )] # Merge layer self . flows += [ Merge ( split_mode )] def forward ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for flow in self . flows : z , log_det = flow ( z ) log_det_tot += log_det return z , log_det_tot def inverse ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_det_tot += log_det return z , log_det_tot __init__ ( param_map , scale = True , scale_map = 'exp' , split_mode = 'channel' ) Constructor Parameters: Name Type Description Default param_map Maps features to shift and scale parameter (if applicable) required scale Flag whether scale shall be applied True scale_map Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow 'exp' split_mode Splitting mode, for possible values see Split class 'channel' Source code in normflows/flows/affine/coupling.py 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 def __init__ ( self , param_map , scale = True , scale_map = \"exp\" , split_mode = \"channel\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Split layer self . flows += [ Split ( split_mode )] # Affine coupling layer self . flows += [ AffineCoupling ( param_map , scale , scale_map )] # Merge layer self . flows += [ Merge ( split_mode )] CCAffineConst Bases: Flow Affine constant flow layer with class-conditional parameters Source code in normflows/flows/affine/coupling.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 class CCAffineConst ( Flow ): \"\"\" Affine constant flow layer with class-conditional parameters \"\"\" def __init__ ( self , shape , num_classes ): super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) self . shape = shape self . s = nn . Parameter ( torch . zeros ( shape )[ None ]) self . t = nn . Parameter ( torch . zeros ( shape )[ None ]) self . s_cc = nn . Parameter ( torch . zeros ( num_classes , np . prod ( shape ))) self . t_cc = nn . Parameter ( torch . zeros ( num_classes , np . prod ( shape ))) self . n_dim = self . s . dim () self . batch_dims = torch . nonzero ( torch . tensor ( self . s . shape ) == 1 , as_tuple = False )[:, 0 ] . tolist () def forward ( self , z , y ): s = self . s + ( y @ self . s_cc ) . view ( - 1 , * self . shape ) t = self . t + ( y @ self . t_cc ) . view ( - 1 , * self . shape ) z_ = z * torch . exp ( s ) + t if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = prod_batch_dims * torch . sum ( s , dim = list ( range ( 1 , self . n_dim ))) return z_ , log_det def inverse ( self , z , y ): s = self . s + ( y @ self . s_cc ) . view ( - 1 , * self . shape ) t = self . t + ( y @ self . t_cc ) . view ( - 1 , * self . shape ) z_ = ( z - t ) * torch . exp ( - s ) if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = - prod_batch_dims * torch . sum ( s , dim = list ( range ( 1 , self . n_dim ))) return z_ , log_det MaskedAffineFlow Bases: Flow RealNVP as introduced in arXiv: 1605.08803 Masked affine flow: f(z) = b * z + (1 - b) * (z * exp(s(b * z)) + t) class AffineHalfFlow(Flow): is MaskedAffineFlow with alternating bit mask NICE is AffineFlow with only shifts (volume preserving) Source code in normflows/flows/affine/coupling.py 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 class MaskedAffineFlow ( Flow ): \"\"\"RealNVP as introduced in [arXiv: 1605.08803](https://arxiv.org/abs/1605.08803) Masked affine flow: ``` f(z) = b * z + (1 - b) * (z * exp(s(b * z)) + t) ``` - class AffineHalfFlow(Flow): is MaskedAffineFlow with alternating bit mask - NICE is AffineFlow with only shifts (volume preserving) \"\"\" def __init__ ( self , b , t = None , s = None ): \"\"\"Constructor Args: b: mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s t: translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied s: scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied \"\"\" super () . __init__ () self . b_cpu = b . view ( 1 , * b . size ()) self . register_buffer ( \"b\" , self . b_cpu ) if s is None : self . s = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"s\" , s ) if t is None : self . t = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"t\" , t ) def forward ( self , z ): z_masked = self . b * z scale = self . s ( z_masked ) nan = torch . tensor ( np . nan , dtype = z . dtype , device = z . device ) scale = torch . where ( torch . isfinite ( scale ), scale , nan ) trans = self . t ( z_masked ) trans = torch . where ( torch . isfinite ( trans ), trans , nan ) z_ = z_masked + ( 1 - self . b ) * ( z * torch . exp ( scale ) + trans ) log_det = torch . sum (( 1 - self . b ) * scale , dim = list ( range ( 1 , self . b . dim ()))) return z_ , log_det def inverse ( self , z ): z_masked = self . b * z scale = self . s ( z_masked ) nan = torch . tensor ( np . nan , dtype = z . dtype , device = z . device ) scale = torch . where ( torch . isfinite ( scale ), scale , nan ) trans = self . t ( z_masked ) trans = torch . where ( torch . isfinite ( trans ), trans , nan ) z_ = z_masked + ( 1 - self . b ) * ( z - trans ) * torch . exp ( - scale ) log_det = - torch . sum (( 1 - self . b ) * scale , dim = list ( range ( 1 , self . b . dim ()))) return z_ , log_det __init__ ( b , t = None , s = None ) Constructor Parameters: Name Type Description Default b mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s required t translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied None s scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied None Source code in normflows/flows/affine/coupling.py 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 def __init__ ( self , b , t = None , s = None ): \"\"\"Constructor Args: b: mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s t: translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied s: scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied \"\"\" super () . __init__ () self . b_cpu = b . view ( 1 , * b . size ()) self . register_buffer ( \"b\" , self . b_cpu ) if s is None : self . s = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"s\" , s ) if t is None : self . t = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"t\" , t ) coupling_test glow GlowBlock Bases: Flow Glow: Generative Flow with Invertible 1\u00d71 Convolutions, arXiv: 1807.03039 One Block of the Glow model, comprised of MaskedAffineFlow (affine coupling layer) Invertible1x1Conv (dropped if there is only one channel) ActNorm (first batch used for initialization) Source code in normflows/flows/affine/glow.py 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 class GlowBlock ( Flow ): \"\"\"Glow: Generative Flow with Invertible 1\u00d71 Convolutions, [arXiv: 1807.03039](https://arxiv.org/abs/1807.03039) One Block of the Glow model, comprised of - MaskedAffineFlow (affine coupling layer) - Invertible1x1Conv (dropped if there is only one channel) - ActNorm (first batch used for initialization) \"\"\" def __init__ ( self , channels , hidden_channels , scale = True , scale_map = \"sigmoid\" , split_mode = \"channel\" , leaky = 0.0 , init_zeros = True , use_lu = True , net_actnorm = False , ): \"\"\"Constructor Args: channels: Number of channels of the data hidden_channels: number of channels in the hidden layer of the ConvNet scale: Flag, whether to include scale in affine coupling layer scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class leaky: Leaky parameter of LeakyReLUs of ConvNet2d init_zeros: Flag whether to initialize last conv layer with zeros use_lu: Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers logscale_factor: Factor which can be used to control the scale of the log scale factor, see [source](https://github.com/openai/glow) \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Coupling layer kernel_size = ( 3 , 1 , 3 ) num_param = 2 if scale else 1 if \"channel\" == split_mode : channels_ = (( channels + 1 ) // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * ( channels // 2 ),) elif \"channel_inv\" == split_mode : channels_ = ( channels // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * (( channels + 1 ) // 2 ),) elif \"checkerboard\" in split_mode : channels_ = ( channels ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * channels ,) else : raise NotImplementedError ( \"Mode \" + split_mode + \" is not implemented.\" ) param_map = nets . ConvNet2d ( channels_ , kernel_size , leaky , init_zeros , actnorm = net_actnorm ) self . flows += [ AffineCouplingBlock ( param_map , scale , scale_map , split_mode )] # Invertible 1x1 convolution if channels > 1 : self . flows += [ Invertible1x1Conv ( channels , use_lu )] # Activation normalization self . flows += [ ActNorm (( channels ,) + ( 1 , 1 ))] def forward ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for flow in self . flows : z , log_det = flow ( z ) log_det_tot += log_det return z , log_det_tot def inverse ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_det_tot += log_det return z , log_det_tot __init__ ( channels , hidden_channels , scale = True , scale_map = 'sigmoid' , split_mode = 'channel' , leaky = 0.0 , init_zeros = True , use_lu = True , net_actnorm = False ) Constructor Parameters: Name Type Description Default channels Number of channels of the data required hidden_channels number of channels in the hidden layer of the ConvNet required scale Flag, whether to include scale in affine coupling layer True scale_map Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow 'sigmoid' split_mode Splitting mode, for possible values see Split class 'channel' leaky Leaky parameter of LeakyReLUs of ConvNet2d 0.0 init_zeros Flag whether to initialize last conv layer with zeros True use_lu Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers True logscale_factor Factor which can be used to control the scale of the log scale factor, see source required Source code in normflows/flows/affine/glow.py 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 def __init__ ( self , channels , hidden_channels , scale = True , scale_map = \"sigmoid\" , split_mode = \"channel\" , leaky = 0.0 , init_zeros = True , use_lu = True , net_actnorm = False , ): \"\"\"Constructor Args: channels: Number of channels of the data hidden_channels: number of channels in the hidden layer of the ConvNet scale: Flag, whether to include scale in affine coupling layer scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class leaky: Leaky parameter of LeakyReLUs of ConvNet2d init_zeros: Flag whether to initialize last conv layer with zeros use_lu: Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers logscale_factor: Factor which can be used to control the scale of the log scale factor, see [source](https://github.com/openai/glow) \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Coupling layer kernel_size = ( 3 , 1 , 3 ) num_param = 2 if scale else 1 if \"channel\" == split_mode : channels_ = (( channels + 1 ) // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * ( channels // 2 ),) elif \"channel_inv\" == split_mode : channels_ = ( channels // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * (( channels + 1 ) // 2 ),) elif \"checkerboard\" in split_mode : channels_ = ( channels ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * channels ,) else : raise NotImplementedError ( \"Mode \" + split_mode + \" is not implemented.\" ) param_map = nets . ConvNet2d ( channels_ , kernel_size , leaky , init_zeros , actnorm = net_actnorm ) self . flows += [ AffineCouplingBlock ( param_map , scale , scale_map , split_mode )] # Invertible 1x1 convolution if channels > 1 : self . flows += [ Invertible1x1Conv ( channels , use_lu )] # Activation normalization self . flows += [ ActNorm (( channels ,) + ( 1 , 1 ))] glow_test base Composite Bases: Flow Composes several flows into one, in the order they are given. Source code in normflows/flows/base.py 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 class Composite ( Flow ): \"\"\" Composes several flows into one, in the order they are given. \"\"\" def __init__ ( self , flows ): \"\"\"Constructor Args: flows: Iterable of flows to composite \"\"\" super () . __init__ () self . _flows = nn . ModuleList ( flows ) @staticmethod def _cascade ( inputs , funcs ): batch_size = inputs . shape [ 0 ] outputs = inputs total_logabsdet = torch . zeros ( batch_size ) for func in funcs : outputs , logabsdet = func ( outputs ) total_logabsdet += logabsdet return outputs , total_logabsdet def forward ( self , inputs ): funcs = self . _flows return self . _cascade ( inputs , funcs ) def inverse ( self , inputs ): funcs = ( flow . inverse for flow in self . _flows [:: - 1 ]) return self . _cascade ( inputs , funcs ) __init__ ( flows ) Constructor Parameters: Name Type Description Default flows Iterable of flows to composite required Source code in normflows/flows/base.py 53 54 55 56 57 58 59 60 def __init__ ( self , flows ): \"\"\"Constructor Args: flows: Iterable of flows to composite \"\"\" super () . __init__ () self . _flows = nn . ModuleList ( flows ) Flow Bases: nn . Module Generic class for flow functions Source code in normflows/flows/base.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 class Flow ( nn . Module ): \"\"\" Generic class for flow functions \"\"\" def __init__ ( self ): super () . __init__ () def forward ( self , z ): \"\"\" Args: z: input variable, first dimension is batch dim Returns: transformed z and log of absolute determinant \"\"\" raise NotImplementedError ( \"Forward pass has not been implemented.\" ) def inverse ( self , z ): raise NotImplementedError ( \"This flow has no algebraic inverse.\" ) forward ( z ) Parameters: Name Type Description Default z input variable, first dimension is batch dim required Returns: Type Description transformed z and log of absolute determinant Source code in normflows/flows/base.py 13 14 15 16 17 18 19 20 21 def forward ( self , z ): \"\"\" Args: z: input variable, first dimension is batch dim Returns: transformed z and log of absolute determinant \"\"\" raise NotImplementedError ( \"Forward pass has not been implemented.\" ) Reverse Bases: Flow Switches the forward transform of a flow layer with its inverse and vice versa Source code in normflows/flows/base.py 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 class Reverse ( Flow ): \"\"\" Switches the forward transform of a flow layer with its inverse and vice versa \"\"\" def __init__ ( self , flow ): \"\"\"Constructor Args: flow: Flow layer to be reversed \"\"\" super () . __init__ () self . flow = flow def forward ( self , z ): return self . flow . inverse ( z ) def inverse ( self , z ): return self . flow . forward ( z ) __init__ ( flow ) Constructor Parameters: Name Type Description Default flow Flow layer to be reversed required Source code in normflows/flows/base.py 32 33 34 35 36 37 38 39 def __init__ ( self , flow ): \"\"\"Constructor Args: flow: Flow layer to be reversed \"\"\" super () . __init__ () self . flow = flow base_test flow_test FlowTest Bases: unittest . TestCase Generic test case for flow modules Source code in normflows/flows/flow_test.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 class FlowTest ( unittest . TestCase ): \"\"\" Generic test case for flow modules \"\"\" def assertClose ( self , actual , expected , atol = None , rtol = None ): assert_close ( actual , expected , atol = atol , rtol = rtol ) def checkForward ( self , flow , inputs ): # Do forward transform outputs , log_det = flow ( inputs ) # Check type assert outputs . dtype == inputs . dtype # Check shape assert outputs . shape == inputs . shape # Return results return outputs , log_det def checkInverse ( self , flow , inputs ): # Do inverse transform outputs , log_det = flow . inverse ( inputs ) # Check type assert outputs . dtype == inputs . dtype # Check shape assert outputs . shape == inputs . shape # Return results return outputs , log_det def checkForwardInverse ( self , flow , inputs , atol = None , rtol = None ): # Check forward outputs , log_det = self . checkForward ( flow , inputs ) # Check inverse input_ , log_det_ = self . checkInverse ( flow , outputs ) # Check identity self . assertClose ( input_ , inputs , atol , rtol ) ld_id = log_det + log_det_ self . assertClose ( ld_id , torch . zeros_like ( ld_id ), atol , rtol ) mixing Invertible1x1Conv Bases: Flow Invertible 1x1 convolution introduced in the Glow paper Assumes 4d input/output tensors of the form NCHW Source code in normflows/flows/mixing.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 class Invertible1x1Conv ( Flow ): \"\"\" Invertible 1x1 convolution introduced in the Glow paper Assumes 4d input/output tensors of the form NCHW \"\"\" def __init__ ( self , num_channels , use_lu = False ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q ) def _assemble_W ( self , inverse = False ): # assemble W from its components (P, L, U, S) L = torch . tril ( self . L , diagonal =- 1 ) + self . eye U = torch . triu ( self . U , diagonal = 1 ) + torch . diag ( self . sign_S * torch . exp ( self . log_S ) ) if inverse : if self . log_S . dtype == torch . float64 : L_inv = torch . inverse ( L ) U_inv = torch . inverse ( U ) else : L_inv = torch . inverse ( L . double ()) . type ( self . log_S . dtype ) U_inv = torch . inverse ( U . double ()) . type ( self . log_S . dtype ) W = U_inv @ L_inv @ self . P . t () else : W = self . P @ L @ U return W def forward ( self , z ): if self . use_lu : W = self . _assemble_W ( inverse = True ) log_det = - torch . sum ( self . log_S ) else : W_dtype = self . W . dtype if W_dtype == torch . float64 : W = torch . inverse ( self . W ) else : W = torch . inverse ( self . W . double ()) . type ( W_dtype ) W = W . view ( * W . size (), 1 , 1 ) log_det = - torch . slogdet ( self . W )[ 1 ] W = W . view ( self . num_channels , self . num_channels , 1 , 1 ) z_ = torch . nn . functional . conv2d ( z , W ) log_det = log_det * z . size ( 2 ) * z . size ( 3 ) return z_ , log_det def inverse ( self , z ): if self . use_lu : W = self . _assemble_W () log_det = torch . sum ( self . log_S ) else : W = self . W log_det = torch . slogdet ( self . W )[ 1 ] W = W . view ( self . num_channels , self . num_channels , 1 , 1 ) z_ = torch . nn . functional . conv2d ( z , W ) log_det = log_det * z . size ( 2 ) * z . size ( 3 ) return z_ , log_det __init__ ( num_channels , use_lu = False ) Constructor Parameters: Name Type Description Default num_channels Number of channels of the data required use_lu Flag whether to parametrize weights through the LU decomposition False Source code in normflows/flows/mixing.py 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 def __init__ ( self , num_channels , use_lu = False ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q ) InvertibleAffine Bases: Flow Invertible affine transformation without shift, i.e. one-dimensional version of the invertible 1x1 convolutions Source code in normflows/flows/mixing.py 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 class InvertibleAffine ( Flow ): \"\"\" Invertible affine transformation without shift, i.e. one-dimensional version of the invertible 1x1 convolutions \"\"\" def __init__ ( self , num_channels , use_lu = True ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q ) def _assemble_W ( self , inverse = False ): # assemble W from its components (P, L, U, S) L = torch . tril ( self . L , diagonal =- 1 ) + self . eye U = torch . triu ( self . U , diagonal = 1 ) + torch . diag ( self . sign_S * torch . exp ( self . log_S ) ) if inverse : if self . log_S . dtype == torch . float64 : L_inv = torch . inverse ( L ) U_inv = torch . inverse ( U ) else : L_inv = torch . inverse ( L . double ()) . type ( self . log_S . dtype ) U_inv = torch . inverse ( U . double ()) . type ( self . log_S . dtype ) W = U_inv @ L_inv @ self . P . t () else : W = self . P @ L @ U return W def forward ( self , z ): if self . use_lu : W = self . _assemble_W ( inverse = True ) log_det = - torch . sum ( self . log_S ) else : W_dtype = self . W . dtype if W_dtype == torch . float64 : W = torch . inverse ( self . W ) else : W = torch . inverse ( self . W . double ()) . type ( W_dtype ) log_det = - torch . slogdet ( self . W )[ 1 ] z_ = z @ W return z_ , log_det def inverse ( self , z ): if self . use_lu : W = self . _assemble_W () log_det = torch . sum ( self . log_S ) else : W = self . W log_det = torch . slogdet ( self . W )[ 1 ] z_ = z @ W return z_ , log_det __init__ ( num_channels , use_lu = True ) Constructor Parameters: Name Type Description Default num_channels Number of channels of the data required use_lu Flag whether to parametrize weights through the LU decomposition True Source code in normflows/flows/mixing.py 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 def __init__ ( self , num_channels , use_lu = True ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q ) LULinearPermute Bases: Flow Fixed permutation combined with a linear transformation parametrized using the LU decomposition, used in https://arxiv.org/abs/1906.04032 Source code in normflows/flows/mixing.py 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 class LULinearPermute ( Flow ): \"\"\" Fixed permutation combined with a linear transformation parametrized using the LU decomposition, used in https://arxiv.org/abs/1906.04032 \"\"\" def __init__ ( self , num_channels , identity_init = True ): \"\"\"Constructor Args: num_channels: Number of dimensions of the data identity_init: Flag, whether to initialize linear transform as identity matrix \"\"\" # Initialize super () . __init__ () # Define modules self . permutation = _RandomPermutation ( num_channels ) self . linear = _LULinear ( num_channels , identity_init = identity_init ) def forward ( self , z ): z , log_det = self . linear . inverse ( z ) z , _ = self . permutation . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , _ = self . permutation ( z ) z , log_det = self . linear ( z ) return z , log_det . view ( - 1 ) __init__ ( num_channels , identity_init = True ) Constructor Parameters: Name Type Description Default num_channels Number of dimensions of the data required identity_init Flag, whether to initialize linear transform as identity matrix True Source code in normflows/flows/mixing.py 541 542 543 544 545 546 547 548 549 550 551 552 553 def __init__ ( self , num_channels , identity_init = True ): \"\"\"Constructor Args: num_channels: Number of dimensions of the data identity_init: Flag, whether to initialize linear transform as identity matrix \"\"\" # Initialize super () . __init__ () # Define modules self . permutation = _RandomPermutation ( num_channels ) self . linear = _LULinear ( num_channels , identity_init = identity_init ) Permute Bases: Flow Permutation features along the channel dimension Source code in normflows/flows/mixing.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 class Permute ( Flow ): \"\"\" Permutation features along the channel dimension \"\"\" def __init__ ( self , num_channels , mode = \"shuffle\" ): \"\"\"Constructor Args: num_channel: Number of channels mode: Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part \"\"\" super () . __init__ () self . mode = mode self . num_channels = num_channels if self . mode == \"shuffle\" : perm = torch . randperm ( self . num_channels ) inv_perm = torch . empty_like ( perm ) . scatter_ ( dim = 0 , index = perm , src = torch . arange ( self . num_channels ) ) self . register_buffer ( \"perm\" , perm ) self . register_buffer ( \"inv_perm\" , inv_perm ) def forward ( self , z ): if self . mode == \"shuffle\" : z = z [:, self . perm , ... ] elif self . mode == \"swap\" : z1 = z [:, : self . num_channels // 2 , ... ] z2 = z [:, self . num_channels // 2 :, ... ] z = torch . cat ([ z2 , z1 ], dim = 1 ) else : raise NotImplementedError ( \"The mode \" + self . mode + \" is not implemented.\" ) log_det = torch . zeros ( len ( z ), device = z . device ) return z , log_det def inverse ( self , z ): if self . mode == \"shuffle\" : z = z [:, self . inv_perm , ... ] elif self . mode == \"swap\" : z1 = z [:, : ( self . num_channels + 1 ) // 2 , ... ] z2 = z [:, ( self . num_channels + 1 ) // 2 :, ... ] z = torch . cat ([ z2 , z1 ], dim = 1 ) else : raise NotImplementedError ( \"The mode \" + self . mode + \" is not implemented.\" ) log_det = torch . zeros ( len ( z ), device = z . device ) return z , log_det __init__ ( num_channels , mode = 'shuffle' ) Constructor Parameters: Name Type Description Default num_channel Number of channels required mode Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part 'shuffle' Source code in normflows/flows/mixing.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 def __init__ ( self , num_channels , mode = \"shuffle\" ): \"\"\"Constructor Args: num_channel: Number of channels mode: Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part \"\"\" super () . __init__ () self . mode = mode self . num_channels = num_channels if self . mode == \"shuffle\" : perm = torch . randperm ( self . num_channels ) inv_perm = torch . empty_like ( perm ) . scatter_ ( dim = 0 , index = perm , src = torch . arange ( self . num_channels ) ) self . register_buffer ( \"perm\" , perm ) self . register_buffer ( \"inv_perm\" , inv_perm ) mixing_test neural_spline autoregressive Implementations of autoregressive transforms. Code taken from https://github.com/bayesiains/nsf autoregressive_test Tests for the autoregressive transforms. Code partially taken from https://github.com/bayesiains/nsf coupling Implementations of various coupling layers. Code taken from https://github.com/bayesiains/nsf Coupling Bases: Flow A base class for coupling layers. Supports 2D inputs (NxD), as well as 4D inputs for images (NxCxHxW). For images the splitting is done on the channel dimension, using the provided 1D mask. Source code in normflows/flows/neural_spline/coupling.py 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 class Coupling ( Flow ): \"\"\"A base class for coupling layers. Supports 2D inputs (NxD), as well as 4D inputs for images (NxCxHxW). For images the splitting is done on the channel dimension, using the provided 1D mask.\"\"\" def __init__ ( self , mask , transform_net_create_fn , unconditional_transform = None ): \"\"\"Constructor. mask: a 1-dim tensor, tuple or list. It indexes inputs as follows: - if `mask[i] > 0`, `input[i]` will be transformed. - if `mask[i] <= 0`, `input[i]` will be passed unchanged. Args: mask \"\"\" mask = torch . as_tensor ( mask ) if mask . dim () != 1 : raise ValueError ( \"Mask must be a 1-dim tensor.\" ) if mask . numel () <= 0 : raise ValueError ( \"Mask can't be empty.\" ) super () . __init__ () self . features = len ( mask ) features_vector = torch . arange ( self . features ) self . register_buffer ( \"identity_features\" , features_vector . masked_select ( mask <= 0 ) ) self . register_buffer ( \"transform_features\" , features_vector . masked_select ( mask > 0 ) ) assert self . num_identity_features + self . num_transform_features == self . features self . transform_net = transform_net_create_fn ( self . num_identity_features , self . num_transform_features * self . _transform_dim_multiplier (), ) if unconditional_transform is None : self . unconditional_transform = None else : self . unconditional_transform = unconditional_transform ( features = self . num_identity_features ) @property def num_identity_features ( self ): return len ( self . identity_features ) @property def num_transform_features ( self ): return len ( self . transform_features ) def forward ( self , inputs , context = None ): if inputs . dim () not in [ 2 , 4 ]: raise ValueError ( \"Inputs must be a 2D or a 4D tensor.\" ) if inputs . shape [ 1 ] != self . features : raise ValueError ( \"Expected features = {} , got {} .\" . format ( self . features , inputs . shape [ 1 ]) ) identity_split = inputs [:, self . identity_features , ... ] transform_split = inputs [:, self . transform_features , ... ] transform_params = self . transform_net ( identity_split , context ) transform_split , logabsdet = self . _coupling_transform_forward ( inputs = transform_split , transform_params = transform_params ) if self . unconditional_transform is not None : identity_split , logabsdet_identity = self . unconditional_transform ( identity_split , context ) logabsdet += logabsdet_identity outputs = torch . empty_like ( inputs ) outputs [:, self . identity_features , ... ] = identity_split outputs [:, self . transform_features , ... ] = transform_split return outputs , logabsdet def inverse ( self , inputs , context = None ): if inputs . dim () not in [ 2 , 4 ]: raise ValueError ( \"Inputs must be a 2D or a 4D tensor.\" ) if inputs . shape [ 1 ] != self . features : raise ValueError ( \"Expected features = {} , got {} .\" . format ( self . features , inputs . shape [ 1 ]) ) identity_split = inputs [:, self . identity_features , ... ] transform_split = inputs [:, self . transform_features , ... ] logabsdet = 0.0 if self . unconditional_transform is not None : identity_split , logabsdet = self . unconditional_transform . inverse ( identity_split , context ) transform_params = self . transform_net ( identity_split , context ) transform_split , logabsdet_split = self . _coupling_transform_inverse ( inputs = transform_split , transform_params = transform_params ) logabsdet += logabsdet_split outputs = torch . empty_like ( inputs ) outputs [:, self . identity_features ] = identity_split outputs [:, self . transform_features ] = transform_split return outputs , logabsdet def _transform_dim_multiplier ( self ): \"\"\"Number of features to output for each transform dimension.\"\"\" raise NotImplementedError () def _coupling_transform_forward ( self , inputs , transform_params ): \"\"\"Forward pass of the coupling transform.\"\"\" raise NotImplementedError () def _coupling_transform_inverse ( self , inputs , transform_params ): \"\"\"Inverse of the coupling transform.\"\"\" raise NotImplementedError () __init__ ( mask , transform_net_create_fn , unconditional_transform = None ) Constructor. mask: a 1-dim tensor, tuple or list. It indexes inputs as follows: if mask[i] > 0 , input[i] will be transformed. if mask[i] <= 0 , input[i] will be passed unchanged. Source code in normflows/flows/neural_spline/coupling.py 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 def __init__ ( self , mask , transform_net_create_fn , unconditional_transform = None ): \"\"\"Constructor. mask: a 1-dim tensor, tuple or list. It indexes inputs as follows: - if `mask[i] > 0`, `input[i]` will be transformed. - if `mask[i] <= 0`, `input[i]` will be passed unchanged. Args: mask \"\"\" mask = torch . as_tensor ( mask ) if mask . dim () != 1 : raise ValueError ( \"Mask must be a 1-dim tensor.\" ) if mask . numel () <= 0 : raise ValueError ( \"Mask can't be empty.\" ) super () . __init__ () self . features = len ( mask ) features_vector = torch . arange ( self . features ) self . register_buffer ( \"identity_features\" , features_vector . masked_select ( mask <= 0 ) ) self . register_buffer ( \"transform_features\" , features_vector . masked_select ( mask > 0 ) ) assert self . num_identity_features + self . num_transform_features == self . features self . transform_net = transform_net_create_fn ( self . num_identity_features , self . num_transform_features * self . _transform_dim_multiplier (), ) if unconditional_transform is None : self . unconditional_transform = None else : self . unconditional_transform = unconditional_transform ( features = self . num_identity_features ) coupling_test Tests for the coupling Transforms. Code partially taken from https://github.com/bayesiains/nsf wrapper AutoregressiveRationalQuadraticSpline Bases: Flow Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see sources Source code in normflows/flows/neural_spline/wrapper.py 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 class AutoregressiveRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see [sources](https://github.com/bayesiains/nsf) \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = False , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = \"linear\" , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , ) def forward ( self , z ): z , log_det = self . mprqat . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . mprqat ( z ) return z , log_det . view ( - 1 ) __init__ ( num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = False , init_identity = True ) Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required num_bins int Number of bins 8 tail_bound int Bound of the spline tails 3 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 permute_mask bool Flag, permutes the mask of the NN False init_identity bool Flag, initialize transform as identity True Source code in normflows/flows/neural_spline/wrapper.py 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = False , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = \"linear\" , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , ) CircularAutoregressiveRationalQuadraticSpline Bases: Flow Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see sources Source code in normflows/flows/neural_spline/wrapper.py 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 class CircularAutoregressiveRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see [sources](https://github.com/bayesiains/nsf) \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = True , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = tails , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , ) def forward ( self , z ): z , log_det = self . mprqat . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . mprqat ( z ) return z , log_det . view ( - 1 ) __init__ ( num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = True , init_identity = True ) Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required ind_circ Iterable Indices of the circular coordinates required num_bins int Number of bins 8 tail_bound int Bound of the spline tails 3 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 permute_mask bool Flag, permutes the mask of the NN True init_identity bool Flag, initialize transform as identity True Source code in normflows/flows/neural_spline/wrapper.py 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = True , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = tails , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , ) CircularCoupledRationalQuadraticSpline Bases: Flow Neural spline flow coupling layer with circular coordinates Source code in normflows/flows/neural_spline/wrapper.py 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 class CircularCoupledRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer with circular coordinates \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , mask = None , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (float or Iterable): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used mask (torch tensor): Mask to be used, alternating masked generated is None init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () if mask is None : mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ) features_vector = torch . arange ( num_input_channels ) identity_features = features_vector . masked_select ( mask <= 0 ) ind_circ = torch . tensor ( ind_circ ) ind_circ_id = [] for i , id in enumerate ( identity_features ): if id in ind_circ : ind_circ_id += [ i ] if torch . is_tensor ( tail_bound ): scale_pf = np . pi / tail_bound [ ind_circ_id ] else : scale_pf = np . pi / tail_bound def transform_net_create_fn ( in_features , out_features ): if len ( ind_circ_id ) > 0 : pf = PeriodicFeaturesElementwise ( in_features , ind_circ_id , scale_pf ) else : pf = None net = ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , preprocessing = pf , ) if init_identity : torch . nn . init . constant_ ( net . final_layer . weight , 0.0 ) torch . nn . init . constant_ ( net . final_layer . bias , np . log ( np . exp ( 1 - DEFAULT_MIN_DERIVATIVE ) - 1 ) ) return net tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . prqct = PiecewiseRationalQuadraticCoupling ( mask = mask , transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , apply_unconditional_transform = True , ) def forward ( self , z ): z , log_det = self . prqct . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . prqct ( z ) return z , log_det . view ( - 1 ) __init__ ( num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , mask = None , init_identity = True ) Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required ind_circ Iterable Indices of the circular coordinates required num_bins int Number of bins 8 tail_bound float or Iterable Bound of the spline tails 3.0 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 reverse_mask bool Flag whether the reverse mask should be used False mask torch tensor Mask to be used, alternating masked generated is None None init_identity bool Flag, initialize transform as identity True Source code in normflows/flows/neural_spline/wrapper.py 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , mask = None , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (float or Iterable): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used mask (torch tensor): Mask to be used, alternating masked generated is None init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () if mask is None : mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ) features_vector = torch . arange ( num_input_channels ) identity_features = features_vector . masked_select ( mask <= 0 ) ind_circ = torch . tensor ( ind_circ ) ind_circ_id = [] for i , id in enumerate ( identity_features ): if id in ind_circ : ind_circ_id += [ i ] if torch . is_tensor ( tail_bound ): scale_pf = np . pi / tail_bound [ ind_circ_id ] else : scale_pf = np . pi / tail_bound def transform_net_create_fn ( in_features , out_features ): if len ( ind_circ_id ) > 0 : pf = PeriodicFeaturesElementwise ( in_features , ind_circ_id , scale_pf ) else : pf = None net = ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , preprocessing = pf , ) if init_identity : torch . nn . init . constant_ ( net . final_layer . weight , 0.0 ) torch . nn . init . constant_ ( net . final_layer . bias , np . log ( np . exp ( 1 - DEFAULT_MIN_DERIVATIVE ) - 1 ) ) return net tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . prqct = PiecewiseRationalQuadraticCoupling ( mask = mask , transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , apply_unconditional_transform = True , ) CoupledRationalQuadraticSpline Bases: Flow Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see source Source code in normflows/flows/neural_spline/wrapper.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 class CoupledRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see [source](https://github.com/bayesiains/nsf) \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tails = \"linear\" , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tails (str): Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval tail_bound (float): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used \"\"\" super () . __init__ () def transform_net_create_fn ( in_features , out_features ): return ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , ) self . prqct = PiecewiseRationalQuadraticCoupling ( mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ), transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , # Setting True corresponds to equations (4), (5), (6) in the NSF paper: apply_unconditional_transform = True , ) def forward ( self , z ): z , log_det = self . prqct . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . prqct ( z ) return z , log_det . view ( - 1 ) __init__ ( num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tails = 'linear' , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False ) Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required num_bins int Number of bins 8 tails str Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval 'linear' tail_bound float Bound of the spline tails 3.0 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 reverse_mask bool Flag whether the reverse mask should be used False Source code in normflows/flows/neural_spline/wrapper.py 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tails = \"linear\" , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tails (str): Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval tail_bound (float): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used \"\"\" super () . __init__ () def transform_net_create_fn ( in_features , out_features ): return ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , ) self . prqct = PiecewiseRationalQuadraticCoupling ( mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ), transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , # Setting True corresponds to equations (4), (5), (6) in the NSF paper: apply_unconditional_transform = True , ) wrapper_test normalization ActNorm Bases: AffineConstFlow An AffineConstFlow but with a data-dependent initialization, where on the very first batch we clever initialize the s,t so that the output is unit gaussian. As described in Glow paper. Source code in normflows/flows/normalization.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 class ActNorm ( AffineConstFlow ): \"\"\" An AffineConstFlow but with a data-dependent initialization, where on the very first batch we clever initialize the s,t so that the output is unit gaussian. As described in Glow paper. \"\"\" def __init__ ( self , * args , ** kwargs ): super () . __init__ ( * args , ** kwargs ) self . data_dep_init_done_cpu = torch . tensor ( 0.0 ) self . register_buffer ( \"data_dep_init_done\" , self . data_dep_init_done_cpu ) def forward ( self , z ): # first batch is used for initialization, c.f. batchnorm if not self . data_dep_init_done > 0.0 : assert self . s is not None and self . t is not None s_init = - torch . log ( z . std ( dim = self . batch_dims , keepdim = True ) + 1e-6 ) self . s . data = s_init . data self . t . data = ( - z . mean ( dim = self . batch_dims , keepdim = True ) * torch . exp ( self . s ) ) . data self . data_dep_init_done = torch . tensor ( 1.0 ) return super () . forward ( z ) def inverse ( self , z ): # first batch is used for initialization, c.f. batchnorm if not self . data_dep_init_done : assert self . s is not None and self . t is not None s_init = torch . log ( z . std ( dim = self . batch_dims , keepdim = True ) + 1e-6 ) self . s . data = s_init . data self . t . data = z . mean ( dim = self . batch_dims , keepdim = True ) . data self . data_dep_init_done = torch . tensor ( 1.0 ) return super () . inverse ( z ) BatchNorm Bases: Flow Batch Normalization with out considering the derivatives of the batch statistics, see arXiv: 1605.08803 Source code in normflows/flows/normalization.py 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 class BatchNorm ( Flow ): \"\"\" Batch Normalization with out considering the derivatives of the batch statistics, see [arXiv: 1605.08803](https://arxiv.org/abs/1605.08803) \"\"\" def __init__ ( self , eps = 1.0e-10 ): super () . __init__ () self . eps_cpu = torch . tensor ( eps ) self . register_buffer ( \"eps\" , self . eps_cpu ) def forward ( self , z ): \"\"\" Do batch norm over batch and sample dimension \"\"\" mean = torch . mean ( z , dim = 0 , keepdims = True ) std = torch . std ( z , dim = 0 , keepdims = True ) z_ = ( z - mean ) / torch . sqrt ( std ** 2 + self . eps ) log_det = torch . log ( 1 / torch . prod ( torch . sqrt ( std ** 2 + self . eps ))) . repeat ( z . size ()[ 0 ] ) return z_ , log_det forward ( z ) Do batch norm over batch and sample dimension Source code in normflows/flows/normalization.py 52 53 54 55 56 57 58 59 60 61 62 def forward ( self , z ): \"\"\" Do batch norm over batch and sample dimension \"\"\" mean = torch . mean ( z , dim = 0 , keepdims = True ) std = torch . std ( z , dim = 0 , keepdims = True ) z_ = ( z - mean ) / torch . sqrt ( std ** 2 + self . eps ) log_det = torch . log ( 1 / torch . prod ( torch . sqrt ( std ** 2 + self . eps ))) . repeat ( z . size ()[ 0 ] ) return z_ , log_det periodic PeriodicShift Bases: Flow Shift and wrap periodic coordinates Source code in normflows/flows/periodic.py 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 class PeriodicShift ( Flow ): \"\"\" Shift and wrap periodic coordinates \"\"\" def __init__ ( self , ind , bound = 1.0 , shift = 0.0 ): \"\"\"Constructor Args: ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval shift: Tensor, shift to be applied \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound if torch . is_tensor ( shift ): self . register_buffer ( \"shift\" , shift ) else : self . shift = shift def forward ( self , z ): z_ = z . clone () z_ [ ... , self . ind ] = ( torch . remainder ( z_ [ ... , self . ind ] + self . shift + self . bound , 2 * self . bound ) - self . bound ) return z_ , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device ) def inverse ( self , z ): z_ = z . clone () z_ [ ... , self . ind ] = ( torch . remainder ( z_ [ ... , self . ind ] - self . shift + self . bound , 2 * self . bound ) - self . bound ) return z_ , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device ) __init__ ( ind , bound = 1.0 , shift = 0.0 ) Constructor Parameters: Name Type Description Default ind Iterable, indices of coordinates to be mapped required bound Float or iterable, bound of interval 1.0 shift Tensor, shift to be applied 0.0 Source code in normflows/flows/periodic.py 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 def __init__ ( self , ind , bound = 1.0 , shift = 0.0 ): \"\"\"Constructor Args: ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval shift: Tensor, shift to be applied \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound if torch . is_tensor ( shift ): self . register_buffer ( \"shift\" , shift ) else : self . shift = shift PeriodicWrap Bases: Flow Map periodic coordinates to fixed interval Source code in normflows/flows/periodic.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 class PeriodicWrap ( Flow ): \"\"\" Map periodic coordinates to fixed interval \"\"\" def __init__ ( self , ind , bound = 1.0 ): \"\"\"Constructor ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound def forward ( self , z ): return z , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device ) def inverse ( self , z ): z_ = z . clone () z_ [ ... , self . ind ] = ( torch . remainder ( z_ [ ... , self . ind ] + self . bound , 2 * self . bound ) - self . bound ) return z_ , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device ) __init__ ( ind , bound = 1.0 ) Constructor ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval Source code in normflows/flows/periodic.py 11 12 13 14 15 16 17 18 19 20 21 22 def __init__ ( self , ind , bound = 1.0 ): \"\"\"Constructor ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound periodic_test planar Planar Bases: Flow Planar flow as introduced in arXiv: 1505.05770 f(z) = z + u * h(w * z + b) Source code in normflows/flows/planar.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 class Planar ( Flow ): \"\"\"Planar flow as introduced in [arXiv: 1505.05770](https://arxiv.org/abs/1505.05770) ``` f(z) = z + u * h(w * z + b) ``` \"\"\" def __init__ ( self , shape , act = \"tanh\" , u = None , w = None , b = None ): \"\"\"Constructor of the planar flow Args: shape: shape of the latent variable z h: nonlinear function h of the planar flow (see definition of f above) u,w,b: optional initialization for parameters \"\"\" super () . __init__ () lim_w = np . sqrt ( 2.0 / np . prod ( shape )) lim_u = np . sqrt ( 2 ) if u is not None : self . u = nn . Parameter ( u ) else : self . u = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . u , - lim_u , lim_u ) if w is not None : self . w = nn . Parameter ( w ) else : self . w = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . w , - lim_w , lim_w ) if b is not None : self . b = nn . Parameter ( b ) else : self . b = nn . Parameter ( torch . zeros ( 1 )) self . act = act if act == \"tanh\" : self . h = torch . tanh elif act == \"leaky_relu\" : self . h = torch . nn . LeakyReLU ( negative_slope = 0.2 ) else : raise NotImplementedError ( \"Nonlinearity is not implemented.\" ) def forward ( self , z ): lin = torch . sum ( self . w * z , list ( range ( 1 , self . w . dim ())), keepdim = True ) + self . b inner = torch . sum ( self . w * self . u ) u = self . u + ( torch . log ( 1 + torch . exp ( inner )) - 1 - inner ) \\ * self . w / torch . sum ( self . w ** 2 ) # constraint w.T * u > -1 if self . act == \"tanh\" : h_ = lambda x : 1 / torch . cosh ( x ) ** 2 elif self . act == \"leaky_relu\" : h_ = lambda x : ( x < 0 ) * ( self . h . negative_slope - 1.0 ) + 1.0 z_ = z + u * self . h ( lin ) log_det = torch . log ( torch . abs ( 1 + torch . sum ( self . w * u ) * h_ ( lin . reshape ( - 1 )))) return z_ , log_det def inverse ( self , z ): if self . act != \"leaky_relu\" : raise NotImplementedError ( \"This flow has no algebraic inverse.\" ) lin = torch . sum ( self . w * z , list ( range ( 1 , self . w . dim ()))) + self . b a = ( lin < 0 ) * ( self . h . negative_slope - 1.0 ) + 1.0 # absorb leakyReLU slope into u inner = torch . sum ( self . w * self . u ) u = self . u + ( torch . log ( 1 + torch . exp ( inner )) - 1 - inner ) \\ * self . w / torch . sum ( self . w ** 2 ) dims = [ - 1 ] + ( u . dim () - 1 ) * [ 1 ] u = a . reshape ( * dims ) * u inner_ = torch . sum ( self . w * u , list ( range ( 1 , self . w . dim ()))) z_ = z - u * ( lin / ( 1 + inner_ )) . reshape ( * dims ) log_det = - torch . log ( torch . abs ( 1 + inner_ )) return z_ , log_det __init__ ( shape , act = 'tanh' , u = None , w = None , b = None ) Constructor of the planar flow Parameters: Name Type Description Default shape shape of the latent variable z required h nonlinear function h of the planar flow (see definition of f above) required u,w,b optional initialization for parameters required Source code in normflows/flows/planar.py 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 def __init__ ( self , shape , act = \"tanh\" , u = None , w = None , b = None ): \"\"\"Constructor of the planar flow Args: shape: shape of the latent variable z h: nonlinear function h of the planar flow (see definition of f above) u,w,b: optional initialization for parameters \"\"\" super () . __init__ () lim_w = np . sqrt ( 2.0 / np . prod ( shape )) lim_u = np . sqrt ( 2 ) if u is not None : self . u = nn . Parameter ( u ) else : self . u = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . u , - lim_u , lim_u ) if w is not None : self . w = nn . Parameter ( w ) else : self . w = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . w , - lim_w , lim_w ) if b is not None : self . b = nn . Parameter ( b ) else : self . b = nn . Parameter ( torch . zeros ( 1 )) self . act = act if act == \"tanh\" : self . h = torch . tanh elif act == \"leaky_relu\" : self . h = torch . nn . LeakyReLU ( negative_slope = 0.2 ) else : raise NotImplementedError ( \"Nonlinearity is not implemented.\" ) planar_test radial Radial Bases: Flow Radial flow as introduced in arXiv: 1505.05770 f(z) = z + beta * h(alpha, r) * (z - z_0) Source code in normflows/flows/radial.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 class Radial ( Flow ): \"\"\"Radial flow as introduced in [arXiv: 1505.05770](https://arxiv.org/abs/1505.05770) ``` f(z) = z + beta * h(alpha, r) * (z - z_0) ``` \"\"\" def __init__ ( self , shape , z_0 = None ): \"\"\"Constructor of the radial flow Args: shape: shape of the latent variable z z_0: parameter of the radial flow \"\"\" super () . __init__ () self . d_cpu = torch . prod ( torch . tensor ( shape )) self . register_buffer ( \"d\" , self . d_cpu ) self . beta = nn . Parameter ( torch . empty ( 1 )) lim = 1.0 / np . prod ( shape ) nn . init . uniform_ ( self . beta , - lim - 1.0 , lim - 1.0 ) self . alpha = nn . Parameter ( torch . empty ( 1 )) nn . init . uniform_ ( self . alpha , - lim , lim ) if z_0 is not None : self . z_0 = nn . Parameter ( z_0 ) else : self . z_0 = nn . Parameter ( torch . randn ( shape )[ None ]) def forward ( self , z ): beta = torch . log ( 1 + torch . exp ( self . beta )) - torch . abs ( self . alpha ) dz = z - self . z_0 r = torch . linalg . vector_norm ( dz , dim = list ( range ( 1 , self . z_0 . dim ())), keepdim = True ) h_arr = beta / ( torch . abs ( self . alpha ) + r ) h_arr_ = - beta * r / ( torch . abs ( self . alpha ) + r ) ** 2 z_ = z + h_arr * dz log_det = ( self . d - 1 ) * torch . log ( 1 + h_arr ) + torch . log ( 1 + h_arr + h_arr_ ) log_det = log_det . reshape ( - 1 ) return z_ , log_det __init__ ( shape , z_0 = None ) Constructor of the radial flow Parameters: Name Type Description Default shape shape of the latent variable z required z_0 parameter of the radial flow None Source code in normflows/flows/radial.py 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 def __init__ ( self , shape , z_0 = None ): \"\"\"Constructor of the radial flow Args: shape: shape of the latent variable z z_0: parameter of the radial flow \"\"\" super () . __init__ () self . d_cpu = torch . prod ( torch . tensor ( shape )) self . register_buffer ( \"d\" , self . d_cpu ) self . beta = nn . Parameter ( torch . empty ( 1 )) lim = 1.0 / np . prod ( shape ) nn . init . uniform_ ( self . beta , - lim - 1.0 , lim - 1.0 ) self . alpha = nn . Parameter ( torch . empty ( 1 )) nn . init . uniform_ ( self . alpha , - lim , lim ) if z_0 is not None : self . z_0 = nn . Parameter ( z_0 ) else : self . z_0 = nn . Parameter ( torch . randn ( shape )[ None ]) radial_test reshape Merge Bases: Split Same as Split but with forward and backward pass interchanged Source code in normflows/flows/reshape.py 88 89 90 91 92 93 94 95 96 97 98 99 100 class Merge ( Split ): \"\"\" Same as Split but with forward and backward pass interchanged \"\"\" def __init__ ( self , mode = \"channel\" ): super () . __init__ ( mode ) def forward ( self , z ): return super () . inverse ( z ) def inverse ( self , z ): return super () . forward ( z ) Split Bases: Flow Split features into two sets Source code in normflows/flows/reshape.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 class Split ( Flow ): \"\"\" Split features into two sets \"\"\" def __init__ ( self , mode = \"channel\" ): \"\"\"Constructor The splitting mode can be: - channel: Splits first feature dimension, usually channels, into two halfs - channel_inv: Same as channel, but with z1 and z2 flipped - checkerboard: Splits features using a checkerboard pattern (last feature dimension must be even) - checkerboard_inv: Same as checkerboard, but with inverted coloring Args: mode: splitting mode \"\"\" super () . __init__ () self . mode = mode def forward ( self , z ): if self . mode == \"channel\" : z1 , z2 = z . chunk ( 2 , dim = 1 ) elif self . mode == \"channel_inv\" : z2 , z1 = z . chunk ( 2 , dim = 1 ) elif \"checkerboard\" in self . mode : n_dims = z . dim () cb0 = 0 cb1 = 1 for i in range ( 1 , n_dims ): cb0_ = cb0 cb1_ = cb1 cb0 = [ cb0_ if j % 2 == 0 else cb1_ for j in range ( z . size ( n_dims - i ))] cb1 = [ cb1_ if j % 2 == 0 else cb0_ for j in range ( z . size ( n_dims - i ))] cb = cb1 if \"inv\" in self . mode else cb0 cb = torch . tensor ( cb )[ None ] . repeat ( len ( z ), * (( n_dims - 1 ) * [ 1 ])) cb = cb . to ( z . device ) z_size = z . size () z1 = z . reshape ( - 1 )[ torch . nonzero ( cb . view ( - 1 ), as_tuple = False )] . view ( * z_size [: - 1 ], - 1 ) z2 = z . reshape ( - 1 )[ torch . nonzero (( 1 - cb ) . view ( - 1 ), as_tuple = False )] . view ( * z_size [: - 1 ], - 1 ) else : raise NotImplementedError ( \"Mode \" + self . mode + \" is not implemented.\" ) log_det = 0 return [ z1 , z2 ], log_det def inverse ( self , z ): z1 , z2 = z if self . mode == \"channel\" : z = torch . cat ([ z1 , z2 ], 1 ) elif self . mode == \"channel_inv\" : z = torch . cat ([ z2 , z1 ], 1 ) elif \"checkerboard\" in self . mode : n_dims = z1 . dim () z_size = list ( z1 . size ()) z_size [ - 1 ] *= 2 cb0 = 0 cb1 = 1 for i in range ( 1 , n_dims ): cb0_ = cb0 cb1_ = cb1 cb0 = [ cb0_ if j % 2 == 0 else cb1_ for j in range ( z_size [ n_dims - i ])] cb1 = [ cb1_ if j % 2 == 0 else cb0_ for j in range ( z_size [ n_dims - i ])] cb = cb1 if \"inv\" in self . mode else cb0 cb = torch . tensor ( cb )[ None ] . repeat ( z_size [ 0 ], * (( n_dims - 1 ) * [ 1 ])) cb = cb . to ( z1 . device ) z1 = z1 [ ... , None ] . repeat ( * ( n_dims * [ 1 ]), 2 ) . view ( * z_size [: - 1 ], - 1 ) z2 = z2 [ ... , None ] . repeat ( * ( n_dims * [ 1 ]), 2 ) . view ( * z_size [: - 1 ], - 1 ) z = cb * z1 + ( 1 - cb ) * z2 else : raise NotImplementedError ( \"Mode \" + self . mode + \" is not implemented.\" ) log_det = 0 return z , log_det __init__ ( mode = 'channel' ) Constructor The splitting mode can be: channel: Splits first feature dimension, usually channels, into two halfs channel_inv: Same as channel, but with z1 and z2 flipped checkerboard: Splits features using a checkerboard pattern (last feature dimension must be even) checkerboard_inv: Same as checkerboard, but with inverted coloring Parameters: Name Type Description Default mode splitting mode 'channel' Source code in normflows/flows/reshape.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 def __init__ ( self , mode = \"channel\" ): \"\"\"Constructor The splitting mode can be: - channel: Splits first feature dimension, usually channels, into two halfs - channel_inv: Same as channel, but with z1 and z2 flipped - checkerboard: Splits features using a checkerboard pattern (last feature dimension must be even) - checkerboard_inv: Same as checkerboard, but with inverted coloring Args: mode: splitting mode \"\"\" super () . __init__ () self . mode = mode Squeeze Bases: Flow Squeeze operation of multi-scale architecture, RealNVP or Glow paper Source code in normflows/flows/reshape.py 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 class Squeeze ( Flow ): \"\"\" Squeeze operation of multi-scale architecture, RealNVP or Glow paper \"\"\" def __init__ ( self ): \"\"\" Constructor \"\"\" super () . __init__ () def forward ( self , z ): log_det = 0 s = z . size () z = z . view ( s [ 0 ], s [ 1 ] // 4 , 2 , 2 , s [ 2 ], s [ 3 ]) z = z . permute ( 0 , 1 , 4 , 2 , 5 , 3 ) . contiguous () z = z . view ( s [ 0 ], s [ 1 ] // 4 , 2 * s [ 2 ], 2 * s [ 3 ]) return z , log_det def inverse ( self , z ): log_det = 0 s = z . size () z = z . view ( * s [: 2 ], s [ 2 ] // 2 , 2 , s [ 3 ] // 2 , 2 ) z = z . permute ( 0 , 1 , 3 , 5 , 2 , 4 ) . contiguous () z = z . view ( s [ 0 ], 4 * s [ 1 ], s [ 2 ] // 2 , s [ 3 ] // 2 ) return z , log_det __init__ () Constructor Source code in normflows/flows/reshape.py 108 109 110 111 112 def __init__ ( self ): \"\"\" Constructor \"\"\" super () . __init__ () residual Residual Bases: Flow Invertible residual net block, wrapper to the implementation of Chen et al., see sources Source code in normflows/flows/residual.py 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 class Residual ( Flow ): \"\"\" Invertible residual net block, wrapper to the implementation of Chen et al., see [sources](https://github.com/rtqichen/residual-flows) \"\"\" def __init__ ( self , net , reverse = True , reduce_memory = True , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" ): \"\"\"Constructor Args: net: Neural network, must be Lipschitz continuous with L < 1 reverse: Flag, if true the map ```f(x) = x + net(x)``` is applied in the inverse pass, otherwise it is done in forward reduce_memory: Flag, if true Neumann series and precomputations, for backward pass in forward pass are done geom_p: Parameter of the geometric distribution used for the Neumann series lamb: Parameter of the geometric distribution used for the Neumann series n_power_series: Number of terms in the Neumann series exact_trace: Flag, if true the trace of the Jacobian is computed exactly brute_force: Flag, if true the Jacobian is computed exactly in 2D n_samples: Number of samples used to estimate power series n_exact_terms: Number of terms always included in the power series n_dist: Distribution used for the power series, either \"geometric\" or \"poisson\" \"\"\" super () . __init__ () self . reverse = reverse self . iresblock = iResBlock ( net , n_samples = n_samples , n_exact_terms = n_exact_terms , neumann_grad = reduce_memory , grad_in_forward = reduce_memory , exact_trace = exact_trace , geom_p = geom_p , lamb = lamb , n_power_series = n_power_series , brute_force = brute_force , n_dist = n_dist , ) def forward ( self , z ): if self . reverse : z , log_det = self . iresblock . inverse ( z , 0 ) else : z , log_det = self . iresblock . forward ( z , 0 ) return z , - log_det . view ( - 1 ) def inverse ( self , z ): if self . reverse : z , log_det = self . iresblock . forward ( z , 0 ) else : z , log_det = self . iresblock . inverse ( z , 0 ) return z , - log_det . view ( - 1 ) __init__ ( net , reverse = True , reduce_memory = True , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = 'geometric' ) Constructor Parameters: Name Type Description Default net Neural network, must be Lipschitz continuous with L < 1 required reverse Flag, if true the map f(x) = x + net(x) is applied in the inverse pass, otherwise it is done in forward True reduce_memory Flag, if true Neumann series and precomputations, for backward pass in forward pass are done True geom_p Parameter of the geometric distribution used for the Neumann series 0.5 lamb Parameter of the geometric distribution used for the Neumann series 2.0 n_power_series Number of terms in the Neumann series None exact_trace Flag, if true the trace of the Jacobian is computed exactly False brute_force Flag, if true the Jacobian is computed exactly in 2D False n_samples Number of samples used to estimate power series 1 n_exact_terms Number of terms always included in the power series 2 n_dist Distribution used for the power series, either \"geometric\" or \"poisson\" 'geometric' Source code in normflows/flows/residual.py 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 def __init__ ( self , net , reverse = True , reduce_memory = True , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" ): \"\"\"Constructor Args: net: Neural network, must be Lipschitz continuous with L < 1 reverse: Flag, if true the map ```f(x) = x + net(x)``` is applied in the inverse pass, otherwise it is done in forward reduce_memory: Flag, if true Neumann series and precomputations, for backward pass in forward pass are done geom_p: Parameter of the geometric distribution used for the Neumann series lamb: Parameter of the geometric distribution used for the Neumann series n_power_series: Number of terms in the Neumann series exact_trace: Flag, if true the trace of the Jacobian is computed exactly brute_force: Flag, if true the Jacobian is computed exactly in 2D n_samples: Number of samples used to estimate power series n_exact_terms: Number of terms always included in the power series n_dist: Distribution used for the power series, either \"geometric\" or \"poisson\" \"\"\" super () . __init__ () self . reverse = reverse self . iresblock = iResBlock ( net , n_samples = n_samples , n_exact_terms = n_exact_terms , neumann_grad = reduce_memory , grad_in_forward = reduce_memory , exact_trace = exact_trace , geom_p = geom_p , lamb = lamb , n_power_series = n_power_series , brute_force = brute_force , n_dist = n_dist , ) iResBlock Bases: nn . Module Source code in normflows/flows/residual.py 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 class iResBlock ( nn . Module ): def __init__ ( self , nnet , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" , neumann_grad = True , grad_in_forward = False , ): \"\"\" Args: nnet: a nn.Module n_power_series: number of power series. If not None, uses a biased approximation to logdet. exact_trace: if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. brute_force: Computes the exact logdet. Only available for 2D inputs. \"\"\" nn . Module . __init__ ( self ) self . nnet = nnet self . n_dist = n_dist self . geom_p = nn . Parameter ( torch . tensor ( np . log ( geom_p ) - np . log ( 1.0 - geom_p ))) self . lamb = nn . Parameter ( torch . tensor ( lamb )) self . n_samples = n_samples self . n_power_series = n_power_series self . exact_trace = exact_trace self . brute_force = brute_force self . n_exact_terms = n_exact_terms self . grad_in_forward = grad_in_forward self . neumann_grad = neumann_grad # store the samples of n. self . register_buffer ( \"last_n_samples\" , torch . zeros ( self . n_samples )) self . register_buffer ( \"last_firmom\" , torch . zeros ( 1 )) self . register_buffer ( \"last_secmom\" , torch . zeros ( 1 )) def forward ( self , x , logpx = None ): if logpx is None : y = x + self . nnet ( x ) return y else : g , logdetgrad = self . _logdetgrad ( x ) return x + g , logpx - logdetgrad def inverse ( self , y , logpy = None ): x = self . _inverse_fixed_point ( y ) if logpy is None : return x else : return x , logpy + self . _logdetgrad ( x )[ 1 ] def _inverse_fixed_point ( self , y , atol = 1e-5 , rtol = 1e-5 ): x , x_prev = y - self . nnet ( y ), y i = 0 tol = atol + y . abs () * rtol while not torch . all (( x - x_prev ) ** 2 / tol < 1 ): x , x_prev = y - self . nnet ( x ), x i += 1 if i > 1000 : break return x def _logdetgrad ( self , x ): \"\"\"Returns g(x) and ```logdet|d(x+g(x))/dx|```\"\"\" with torch . enable_grad (): if ( self . brute_force or not self . training ) and ( x . ndimension () == 2 and x . shape [ 1 ] == 2 ): ########################################### # Brute-force compute Jacobian determinant. ########################################### x = x . requires_grad_ ( True ) g = self . nnet ( x ) # Brute-force logdet only available for 2D. jac = batch_jacobian ( g , x ) batch_dets = ( jac [:, 0 , 0 ] + 1 ) * ( jac [:, 1 , 1 ] + 1 ) - jac [ :, 0 , 1 ] * jac [:, 1 , 0 ] return g , torch . log ( torch . abs ( batch_dets )) . view ( - 1 , 1 ) if self . n_dist == \"geometric\" : geom_p = torch . sigmoid ( self . geom_p ) . item () sample_fn = lambda m : geometric_sample ( geom_p , m ) rcdf_fn = lambda k , offset : geometric_1mcdf ( geom_p , k , offset ) elif self . n_dist == \"poisson\" : lamb = self . lamb . item () sample_fn = lambda m : poisson_sample ( lamb , m ) rcdf_fn = lambda k , offset : poisson_1mcdf ( lamb , k , offset ) if self . training : if self . n_power_series is None : # Unbiased estimation. lamb = self . lamb . item () n_samples = sample_fn ( self . n_samples ) n_power_series = max ( n_samples ) + self . n_exact_terms coeff_fn = ( lambda k : 1 / rcdf_fn ( k , self . n_exact_terms ) * sum ( n_samples >= k - self . n_exact_terms ) / len ( n_samples ) ) else : # Truncated estimation. n_power_series = self . n_power_series coeff_fn = lambda k : 1.0 else : # Unbiased estimation with more exact terms. lamb = self . lamb . item () n_samples = sample_fn ( self . n_samples ) n_power_series = max ( n_samples ) + 20 coeff_fn = ( lambda k : 1 / rcdf_fn ( k , 20 ) * sum ( n_samples >= k - 20 ) / len ( n_samples ) ) if not self . exact_trace : #################################### # Power series with trace estimator. #################################### vareps = torch . randn_like ( x ) # Choose the type of estimator. if self . training and self . neumann_grad : estimator_fn = neumann_logdet_estimator else : estimator_fn = basic_logdet_estimator # Do backprop-in-forward to save memory. if self . training and self . grad_in_forward : g , logdetgrad = mem_eff_wrapper ( estimator_fn , self . nnet , x , n_power_series , vareps , coeff_fn , self . training , ) else : x = x . requires_grad_ ( True ) g = self . nnet ( x ) logdetgrad = estimator_fn ( g , x , n_power_series , vareps , coeff_fn , self . training ) else : ############################################ # Power series with exact trace computation. ############################################ x = x . requires_grad_ ( True ) g = self . nnet ( x ) jac = batch_jacobian ( g , x ) logdetgrad = batch_trace ( jac ) jac_k = jac for k in range ( 2 , n_power_series + 1 ): jac_k = torch . bmm ( jac , jac_k ) logdetgrad = logdetgrad + ( - 1 ) ** ( k + 1 ) / k * coeff_fn ( k ) * batch_trace ( jac_k ) if self . training and self . n_power_series is None : self . last_n_samples . copy_ ( torch . tensor ( n_samples ) . to ( self . last_n_samples ) ) estimator = logdetgrad . detach () self . last_firmom . copy_ ( torch . mean ( estimator ) . to ( self . last_firmom )) self . last_secmom . copy_ ( torch . mean ( estimator ** 2 ) . to ( self . last_secmom )) return g , logdetgrad . view ( - 1 , 1 ) def extra_repr ( self ): return \"dist= {} , n_samples= {} , n_power_series= {} , neumann_grad= {} , exact_trace= {} , brute_force= {} \" . format ( self . n_dist , self . n_samples , self . n_power_series , self . neumann_grad , self . exact_trace , self . brute_force , ) __init__ ( nnet , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = 'geometric' , neumann_grad = True , grad_in_forward = False ) Parameters: Name Type Description Default nnet a nn.Module required n_power_series number of power series. If not None, uses a biased approximation to logdet. None exact_trace if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. False brute_force Computes the exact logdet. Only available for 2D inputs. False Source code in normflows/flows/residual.py 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 def __init__ ( self , nnet , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" , neumann_grad = True , grad_in_forward = False , ): \"\"\" Args: nnet: a nn.Module n_power_series: number of power series. If not None, uses a biased approximation to logdet. exact_trace: if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. brute_force: Computes the exact logdet. Only available for 2D inputs. \"\"\" nn . Module . __init__ ( self ) self . nnet = nnet self . n_dist = n_dist self . geom_p = nn . Parameter ( torch . tensor ( np . log ( geom_p ) - np . log ( 1.0 - geom_p ))) self . lamb = nn . Parameter ( torch . tensor ( lamb )) self . n_samples = n_samples self . n_power_series = n_power_series self . exact_trace = exact_trace self . brute_force = brute_force self . n_exact_terms = n_exact_terms self . grad_in_forward = grad_in_forward self . neumann_grad = neumann_grad # store the samples of n. self . register_buffer ( \"last_n_samples\" , torch . zeros ( self . n_samples )) self . register_buffer ( \"last_firmom\" , torch . zeros ( 1 )) self . register_buffer ( \"last_secmom\" , torch . zeros ( 1 )) residual_test stochastic HamiltonianMonteCarlo Bases: Flow Flow layer using the HMC proposal in Stochastic Normalising Flows See arXiv: 2002.06707 Source code in normflows/flows/stochastic.py 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 class HamiltonianMonteCarlo ( Flow ): \"\"\"Flow layer using the HMC proposal in Stochastic Normalising Flows See [arXiv: 2002.06707](https://arxiv.org/abs/2002.06707) \"\"\" def __init__ ( self , target , steps , log_step_size , log_mass , max_abs_grad = None ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. steps: The number of leapfrog steps log_step_size: The log step size used in the leapfrog integrator. shape (dim) log_mass: The log_mass determining the variance of the momentum samples. shape (dim) max_abs_grad: Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability.\"\"\" super () . __init__ () self . target = target self . steps = steps self . register_parameter ( \"log_step_size\" , torch . nn . Parameter ( log_step_size )) self . register_parameter ( \"log_mass\" , torch . nn . Parameter ( log_mass )) self . max_abs_grad = max_abs_grad def forward ( self , z ): # Draw momentum p = torch . randn_like ( z ) * torch . exp ( 0.5 * self . log_mass ) # leapfrog z_new = z . clone () p_new = p . clone () step_size = torch . exp ( self . log_step_size ) for i in range ( self . steps ): p_half = p_new - ( step_size / 2.0 ) * - self . gradlogP ( z_new ) z_new = z_new + step_size * ( p_half / torch . exp ( self . log_mass )) p_new = p_half - ( step_size / 2.0 ) * - self . gradlogP ( z_new ) # Metropolis Hastings correction probabilities = torch . exp ( self . target . log_prob ( z_new ) - self . target . log_prob ( z ) - 0.5 * torch . sum ( p_new ** 2 / torch . exp ( self . log_mass ), 1 ) + 0.5 * torch . sum ( p ** 2 / torch . exp ( self . log_mass ), 1 ) ) uniforms = torch . rand_like ( probabilities ) mask = uniforms < probabilities z_out = torch . where ( mask . unsqueeze ( 1 ), z_new , z ) return z_out , self . target . log_prob ( z ) - self . target . log_prob ( z_out ) def inverse ( self , z ): return self . forward ( z ) def gradlogP ( self , z ): z_ = z . detach () . requires_grad_ () logp = self . target . log_prob ( z_ ) grad = torch . autograd . grad ( logp , z_ , grad_outputs = torch . ones_like ( logp ))[ 0 ] if self . max_abs_grad : grad = torch . clamp ( grad , max = self . max_abs_grad , min =- self . max_abs_grad ) return grad __init__ ( target , steps , log_step_size , log_mass , max_abs_grad = None ) Constructor Parameters: Name Type Description Default target The stationary distribution of this Markov transition, i.e. the target distribution to sample from. required steps The number of leapfrog steps required log_step_size The log step size used in the leapfrog integrator. shape (dim) required log_mass The log_mass determining the variance of the momentum samples. shape (dim) required max_abs_grad Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability. None Source code in normflows/flows/stochastic.py 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 def __init__ ( self , target , steps , log_step_size , log_mass , max_abs_grad = None ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. steps: The number of leapfrog steps log_step_size: The log step size used in the leapfrog integrator. shape (dim) log_mass: The log_mass determining the variance of the momentum samples. shape (dim) max_abs_grad: Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability.\"\"\" super () . __init__ () self . target = target self . steps = steps self . register_parameter ( \"log_step_size\" , torch . nn . Parameter ( log_step_size )) self . register_parameter ( \"log_mass\" , torch . nn . Parameter ( log_mass )) self . max_abs_grad = max_abs_grad MetropolisHastings Bases: Flow Sampling through Metropolis Hastings in Stochastic Normalizing Flow See arXiv: 2002.06707 Source code in normflows/flows/stochastic.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class MetropolisHastings ( Flow ): \"\"\"Sampling through Metropolis Hastings in Stochastic Normalizing Flow See [arXiv: 2002.06707](https://arxiv.org/abs/2002.06707) \"\"\" def __init__ ( self , target , proposal , steps ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. proposal: Proposal distribution steps: Number of MCMC steps to perform \"\"\" super () . __init__ () self . target = target self . proposal = proposal self . steps = steps def forward ( self , z ): # Initialize number of samples and log(det) num_samples = len ( z ) log_det = torch . zeros ( num_samples , dtype = z . dtype , device = z . device ) # Get log(p) for current samples log_p = self . target . log_prob ( z ) for i in range ( self . steps ): # Make proposal and get log(p) z_ , log_p_diff = self . proposal ( z ) log_p_ = self . target . log_prob ( z_ ) # Make acceptance decision w = torch . rand ( num_samples , dtype = z . dtype , device = z . device ) log_w_accept = log_p_ - log_p + log_p_diff w_accept = torch . clamp ( torch . exp ( log_w_accept ), max = 1 ) accept = w <= w_accept # Update samples, log(det), and log(p) z = torch . where ( accept . unsqueeze ( 1 ), z_ , z ) log_det_ = log_p - log_p_ log_det = torch . where ( accept , log_det + log_det_ , log_det ) log_p = torch . where ( accept , log_p_ , log_p ) return z , log_det def inverse ( self , z ): # Equivalent to forward pass return self . forward ( z ) __init__ ( target , proposal , steps ) Constructor Parameters: Name Type Description Default target The stationary distribution of this Markov transition, i.e. the target distribution to sample from. required proposal Proposal distribution required steps Number of MCMC steps to perform required Source code in normflows/flows/stochastic.py 12 13 14 15 16 17 18 19 20 21 22 23 def __init__ ( self , target , proposal , steps ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. proposal: Proposal distribution steps: Number of MCMC steps to perform \"\"\" super () . __init__ () self . target = target self . proposal = proposal self . steps = steps stochastic_test nets cnn ConvNet2d Bases: nn . Module Convolutional Neural Network with leaky ReLU nonlinearities Source code in normflows/nets/cnn.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 class ConvNet2d ( nn . Module ): \"\"\" Convolutional Neural Network with leaky ReLU nonlinearities \"\"\" def __init__ ( self , channels , kernel_size , leaky = 0.0 , init_zeros = True , actnorm = False , weight_std = None , ): \"\"\"Constructor Args: channels: List of channels of conv layers, first entry is in_channels kernel_size: List of kernel sizes, same for height and width leaky: Leaky part of ReLU init_zeros: Flag whether last layer shall be initialized with zeros scale_output: Flag whether to scale output with a log scale parameter logscale_factor: Constant factor to be multiplied to log scaling actnorm: Flag whether activation normalization shall be done after each conv layer except output weight_std: Fixed std used to initialize every layer \"\"\" super () . __init__ () # Build network net = nn . ModuleList ([]) for i in range ( len ( kernel_size ) - 1 ): conv = nn . Conv2d ( channels [ i ], channels [ i + 1 ], kernel_size [ i ], padding = kernel_size [ i ] // 2 , bias = ( not actnorm ), ) if weight_std is not None : conv . weight . data . normal_ ( mean = 0.0 , std = weight_std ) net . append ( conv ) if actnorm : net . append ( utils . ActNorm (( channels [ i + 1 ],) + ( 1 , 1 ))) net . append ( nn . LeakyReLU ( leaky )) i = len ( kernel_size ) net . append ( nn . Conv2d ( channels [ i - 1 ], channels [ i ], kernel_size [ i - 1 ], padding = kernel_size [ i - 1 ] // 2 , ) ) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) self . net = nn . Sequential ( * net ) def forward ( self , x ): return self . net ( x ) __init__ ( channels , kernel_size , leaky = 0.0 , init_zeros = True , actnorm = False , weight_std = None ) Constructor Parameters: Name Type Description Default channels List of channels of conv layers, first entry is in_channels required kernel_size List of kernel sizes, same for height and width required leaky Leaky part of ReLU 0.0 init_zeros Flag whether last layer shall be initialized with zeros True scale_output Flag whether to scale output with a log scale parameter required logscale_factor Constant factor to be multiplied to log scaling required actnorm Flag whether activation normalization shall be done after each conv layer except output False weight_std Fixed std used to initialize every layer None Source code in normflows/nets/cnn.py 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 def __init__ ( self , channels , kernel_size , leaky = 0.0 , init_zeros = True , actnorm = False , weight_std = None , ): \"\"\"Constructor Args: channels: List of channels of conv layers, first entry is in_channels kernel_size: List of kernel sizes, same for height and width leaky: Leaky part of ReLU init_zeros: Flag whether last layer shall be initialized with zeros scale_output: Flag whether to scale output with a log scale parameter logscale_factor: Constant factor to be multiplied to log scaling actnorm: Flag whether activation normalization shall be done after each conv layer except output weight_std: Fixed std used to initialize every layer \"\"\" super () . __init__ () # Build network net = nn . ModuleList ([]) for i in range ( len ( kernel_size ) - 1 ): conv = nn . Conv2d ( channels [ i ], channels [ i + 1 ], kernel_size [ i ], padding = kernel_size [ i ] // 2 , bias = ( not actnorm ), ) if weight_std is not None : conv . weight . data . normal_ ( mean = 0.0 , std = weight_std ) net . append ( conv ) if actnorm : net . append ( utils . ActNorm (( channels [ i + 1 ],) + ( 1 , 1 ))) net . append ( nn . LeakyReLU ( leaky )) i = len ( kernel_size ) net . append ( nn . Conv2d ( channels [ i - 1 ], channels [ i ], kernel_size [ i - 1 ], padding = kernel_size [ i - 1 ] // 2 , ) ) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) self . net = nn . Sequential ( * net ) lipschitz LipschitzCNN Bases: nn . Module Convolutional neural network which is Lipschitz continuous with Lipschitz constant L < 1 Source code in normflows/nets/lipschitz.py 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 class LipschitzCNN ( nn . Module ): \"\"\" Convolutional neural network which is Lipschitz continuous with Lipschitz constant L < 1 \"\"\" def __init__ ( self , channels , kernel_size , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\"Constructor Args: channels: Integer list with the number of channels of the layers kernel_size: Integer list of kernel sizes of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( kernel_size ) self . channels = channels self . kernel_size = kernel_size self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormConv2d ( in_channels = channels [ i ], out_channels = channels [ i + 1 ], kernel_size = kernel_size [ i ], stride = 1 , padding = kernel_size [ i ] // 2 , bias = True , coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers ) def forward ( self , x ): return self . net ( x ) __init__ ( channels , kernel_size , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True ) Constructor Parameters: Name Type Description Default channels Integer list with the number of channels of the layers required kernel_size Integer list of kernel sizes of the layers required lipschitz_const Maximum Lipschitz constant of each layer 0.97 max_lipschitz_iter Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used 5 lipschitz_tolerance Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 None init_zeros Flag, whether to initialize last layer approximately with zeros True Source code in normflows/nets/lipschitz.py 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 def __init__ ( self , channels , kernel_size , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\"Constructor Args: channels: Integer list with the number of channels of the layers kernel_size: Integer list of kernel sizes of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( kernel_size ) self . channels = channels self . kernel_size = kernel_size self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormConv2d ( in_channels = channels [ i ], out_channels = channels [ i + 1 ], kernel_size = kernel_size [ i ], stride = 1 , padding = kernel_size [ i ] // 2 , bias = True , coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers ) LipschitzMLP Bases: nn . Module Fully connected neural net which is Lipschitz continuou with Lipschitz constant L < 1 Source code in normflows/nets/lipschitz.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 class LipschitzMLP ( nn . Module ): \"\"\"Fully connected neural net which is Lipschitz continuou with Lipschitz constant L < 1\"\"\" def __init__ ( self , channels , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\" Constructor channels: Integer list with the number of channels of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( channels ) - 1 self . channels = channels self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormLinear ( in_features = channels [ i ], out_features = channels [ i + 1 ], coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers ) def forward ( self , x ): return self . net ( x ) __init__ ( channels , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True ) Constructor channels: Integer list with the number of channels of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros Source code in normflows/nets/lipschitz.py 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 def __init__ ( self , channels , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\" Constructor channels: Integer list with the number of channels of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( channels ) - 1 self . channels = channels self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormLinear ( in_features = channels [ i ], out_features = channels [ i + 1 ], coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers ) projmax_ ( v ) Inplace argmax on absolute value. Source code in normflows/nets/lipschitz.py 651 652 653 654 655 656 def projmax_ ( v ): \"\"\"Inplace argmax on absolute value.\"\"\" ind = torch . argmax ( torch . abs ( v )) v . zero_ () v [ ind ] = 1 return v made Implementation of MADE. Code taken from https://github.com/bayesiains/nsf MADE Bases: nn . Module Implementation of MADE. It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no). Source code in normflows/nets/made.py 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 class MADE ( nn . Module ): \"\"\"Implementation of MADE. It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no). \"\"\" def __init__ ( self , features , hidden_features , context_features = None , num_blocks = 2 , output_multiplier = 1 , use_residual_blocks = True , random_mask = False , permute_mask = False , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , preprocessing = None , ): if use_residual_blocks and random_mask : raise ValueError ( \"Residual blocks can't be used with random masks.\" ) super () . __init__ () # Preprocessing if preprocessing is None : self . preprocessing = lambda inputs : inputs else : self . preprocessing = preprocessing # Initial layer. input_degrees_ = _get_input_degrees ( features ) if permute_mask : input_degrees_ = input_degrees_ [ torch . randperm ( features )] self . initial_layer = MaskedLinear ( in_degrees = input_degrees_ , out_features = hidden_features , autoregressive_features = features , random_mask = random_mask , is_output = False , ) if context_features is not None : self . context_layer = nn . Linear ( context_features , hidden_features ) # Residual blocks. blocks = [] if use_residual_blocks : block_constructor = MaskedResidualBlock else : block_constructor = MaskedFeedforwardBlock prev_out_degrees = self . initial_layer . degrees for _ in range ( num_blocks ): blocks . append ( block_constructor ( in_degrees = prev_out_degrees , autoregressive_features = features , context_features = context_features , random_mask = random_mask , activation = activation , dropout_probability = dropout_probability , use_batch_norm = use_batch_norm , ) ) prev_out_degrees = blocks [ - 1 ] . degrees self . blocks = nn . ModuleList ( blocks ) # Final layer. self . final_layer = MaskedLinear ( in_degrees = prev_out_degrees , out_features = features * output_multiplier , autoregressive_features = features , random_mask = random_mask , is_output = True , out_degrees_ = input_degrees_ , ) def forward ( self , inputs , context = None ): outputs = self . preprocessing ( inputs ) outputs = self . initial_layer ( outputs ) if context is not None : outputs += self . context_layer ( context ) for block in self . blocks : outputs = block ( outputs , context ) outputs = self . final_layer ( outputs ) return outputs MaskedFeedforwardBlock Bases: nn . Module A feedforward block based on a masked linear module. NOTE In this implementation, the number of output features is taken to be equal to the number of input features. Source code in normflows/nets/made.py 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 class MaskedFeedforwardBlock ( nn . Module ): \"\"\"A feedforward block based on a masked linear module. **NOTE** In this implementation, the number of output features is taken to be equal to the number of input features. \"\"\" def __init__ ( self , in_degrees , autoregressive_features , context_features = None , random_mask = False , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , ): super () . __init__ () features = len ( in_degrees ) # Batch norm. if use_batch_norm : self . batch_norm = nn . BatchNorm1d ( features , eps = 1e-3 ) else : self . batch_norm = None if context_features is not None : raise NotImplementedError () # Masked linear. self . linear = MaskedLinear ( in_degrees = in_degrees , out_features = features , autoregressive_features = autoregressive_features , random_mask = random_mask , is_output = False , ) self . degrees = self . linear . degrees # Activation and dropout. self . activation = activation self . dropout = nn . Dropout ( p = dropout_probability ) def forward ( self , inputs , context = None ): if context is not None : raise NotImplementedError () if self . batch_norm : outputs = self . batch_norm ( inputs ) else : outputs = inputs outputs = self . linear ( outputs ) outputs = self . activation ( outputs ) outputs = self . dropout ( outputs ) return outputs MaskedLinear Bases: nn . Linear A linear module with a masked weight matrix. Source code in normflows/nets/made.py 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 class MaskedLinear ( nn . Linear ): \"\"\"A linear module with a masked weight matrix.\"\"\" def __init__ ( self , in_degrees , out_features , autoregressive_features , random_mask , is_output , bias = True , out_degrees_ = None , ): super () . __init__ ( in_features = len ( in_degrees ), out_features = out_features , bias = bias ) mask , degrees = self . _get_mask_and_degrees ( in_degrees = in_degrees , out_features = out_features , autoregressive_features = autoregressive_features , random_mask = random_mask , is_output = is_output , out_degrees_ = out_degrees_ , ) self . register_buffer ( \"mask\" , mask ) self . register_buffer ( \"degrees\" , degrees ) @classmethod def _get_mask_and_degrees ( cls , in_degrees , out_features , autoregressive_features , random_mask , is_output , out_degrees_ = None , ): if is_output : if out_degrees_ is None : out_degrees_ = _get_input_degrees ( autoregressive_features ) out_degrees = tile ( out_degrees_ , out_features // autoregressive_features ) mask = ( out_degrees [ ... , None ] > in_degrees ) . float () else : if random_mask : min_in_degree = torch . min ( in_degrees ) . item () min_in_degree = min ( min_in_degree , autoregressive_features - 1 ) out_degrees = torch . randint ( low = min_in_degree , high = autoregressive_features , size = [ out_features ], dtype = torch . long , ) else : max_ = max ( 1 , autoregressive_features - 1 ) min_ = min ( 1 , autoregressive_features - 1 ) out_degrees = torch . arange ( out_features ) % max_ + min_ mask = ( out_degrees [ ... , None ] >= in_degrees ) . float () return mask , out_degrees def forward ( self , x ): return F . linear ( x , self . weight * self . mask , self . bias ) MaskedResidualBlock Bases: nn . Module A residual block containing masked linear modules. Source code in normflows/nets/made.py 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 class MaskedResidualBlock ( nn . Module ): \"\"\"A residual block containing masked linear modules.\"\"\" def __init__ ( self , in_degrees , autoregressive_features , context_features = None , random_mask = False , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , zero_initialization = True , ): if random_mask : raise ValueError ( \"Masked residual block can't be used with random masks.\" ) super () . __init__ () features = len ( in_degrees ) if context_features is not None : self . context_layer = nn . Linear ( context_features , features ) # Batch norm. self . use_batch_norm = use_batch_norm if use_batch_norm : self . batch_norm_layers = nn . ModuleList ( [ nn . BatchNorm1d ( features , eps = 1e-3 ) for _ in range ( 2 )] ) # Masked linear. linear_0 = MaskedLinear ( in_degrees = in_degrees , out_features = features , autoregressive_features = autoregressive_features , random_mask = False , is_output = False , ) linear_1 = MaskedLinear ( in_degrees = linear_0 . degrees , out_features = features , autoregressive_features = autoregressive_features , random_mask = False , is_output = False , ) self . linear_layers = nn . ModuleList ([ linear_0 , linear_1 ]) self . degrees = linear_1 . degrees if torch . all ( self . degrees >= in_degrees ) . item () != 1 : raise RuntimeError ( \"In a masked residual block, the output degrees can't be\" \" less than the corresponding input degrees.\" ) # Activation and dropout self . activation = activation self . dropout = nn . Dropout ( p = dropout_probability ) # Initialization. if zero_initialization : init . uniform_ ( self . linear_layers [ - 1 ] . weight , a =- 1e-3 , b = 1e-3 ) init . uniform_ ( self . linear_layers [ - 1 ] . bias , a =- 1e-3 , b = 1e-3 ) def forward ( self , inputs , context = None ): temps = inputs if self . use_batch_norm : temps = self . batch_norm_layers [ 0 ]( temps ) temps = self . activation ( temps ) temps = self . linear_layers [ 0 ]( temps ) if self . use_batch_norm : temps = self . batch_norm_layers [ 1 ]( temps ) temps = self . activation ( temps ) temps = self . dropout ( temps ) temps = self . linear_layers [ 1 ]( temps ) if context is not None : temps = F . glu ( torch . cat (( temps , self . context_layer ( context )), dim = 1 ), dim = 1 ) return inputs + temps made_test Tests for MADE. Code partially taken from https://github.com/bayesiains/nsf mlp MLP Bases: nn . Module A multilayer perceptron with Leaky ReLU nonlinearities Source code in normflows/nets/mlp.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 class MLP ( nn . Module ): \"\"\" A multilayer perceptron with Leaky ReLU nonlinearities \"\"\" def __init__ ( self , layers , leaky = 0.0 , score_scale = None , output_fn = None , output_scale = None , init_zeros = False , dropout = None , ): \"\"\" layers: list of layer sizes from start to end leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used score_scale: Factor to apply to the scores, i.e. output before output_fn. output_fn: String, function to be applied to the output, either None, \"sigmoid\", \"relu\", \"tanh\", or \"clampexp\" output_scale: Rescale outputs if output_fn is specified, i.e. ```scale * output_fn(out / scale)``` init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see [arXiv 1807.03039](https://arxiv.org/abs/1807.03039)) dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done \"\"\" super () . __init__ () net = nn . ModuleList ([]) for k in range ( len ( layers ) - 2 ): net . append ( nn . Linear ( layers [ k ], layers [ k + 1 ])) net . append ( nn . LeakyReLU ( leaky )) if dropout is not None : net . append ( nn . Dropout ( p = dropout )) net . append ( nn . Linear ( layers [ - 2 ], layers [ - 1 ])) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) if output_fn is not None : if score_scale is not None : net . append ( utils . ConstScaleLayer ( score_scale )) if output_fn == \"sigmoid\" : net . append ( nn . Sigmoid ()) elif output_fn == \"relu\" : net . append ( nn . ReLU ()) elif output_fn == \"tanh\" : net . append ( nn . Tanh ()) elif output_fn == \"clampexp\" : net . append ( utils . ClampExp ()) else : NotImplementedError ( \"This output function is not implemented.\" ) if output_scale is not None : net . append ( utils . ConstScaleLayer ( output_scale )) self . net = nn . Sequential ( * net ) def forward ( self , x ): return self . net ( x ) __init__ ( layers , leaky = 0.0 , score_scale = None , output_fn = None , output_scale = None , init_zeros = False , dropout = None ) layers: list of layer sizes from start to end leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used score_scale: Factor to apply to the scores, i.e. output before output_fn. output_fn: String, function to be applied to the output, either None, \"sigmoid\", \"relu\", \"tanh\", or \"clampexp\" output_scale: Rescale outputs if output_fn is specified, i.e. scale * output_fn(out / scale) init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see arXiv 1807.03039 ) dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done Source code in normflows/nets/mlp.py 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def __init__ ( self , layers , leaky = 0.0 , score_scale = None , output_fn = None , output_scale = None , init_zeros = False , dropout = None , ): \"\"\" layers: list of layer sizes from start to end leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used score_scale: Factor to apply to the scores, i.e. output before output_fn. output_fn: String, function to be applied to the output, either None, \"sigmoid\", \"relu\", \"tanh\", or \"clampexp\" output_scale: Rescale outputs if output_fn is specified, i.e. ```scale * output_fn(out / scale)``` init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see [arXiv 1807.03039](https://arxiv.org/abs/1807.03039)) dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done \"\"\" super () . __init__ () net = nn . ModuleList ([]) for k in range ( len ( layers ) - 2 ): net . append ( nn . Linear ( layers [ k ], layers [ k + 1 ])) net . append ( nn . LeakyReLU ( leaky )) if dropout is not None : net . append ( nn . Dropout ( p = dropout )) net . append ( nn . Linear ( layers [ - 2 ], layers [ - 1 ])) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) if output_fn is not None : if score_scale is not None : net . append ( utils . ConstScaleLayer ( score_scale )) if output_fn == \"sigmoid\" : net . append ( nn . Sigmoid ()) elif output_fn == \"relu\" : net . append ( nn . ReLU ()) elif output_fn == \"tanh\" : net . append ( nn . Tanh ()) elif output_fn == \"clampexp\" : net . append ( utils . ClampExp ()) else : NotImplementedError ( \"This output function is not implemented.\" ) if output_scale is not None : net . append ( utils . ConstScaleLayer ( output_scale )) self . net = nn . Sequential ( * net ) resnet ResidualBlock Bases: nn . Module A general-purpose residual block. Works only with 1-dim inputs. Source code in normflows/nets/resnet.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 class ResidualBlock ( nn . Module ): \"\"\"A general-purpose residual block. Works only with 1-dim inputs.\"\"\" def __init__ ( self , features , context_features , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , zero_initialization = True , ): super () . __init__ () self . activation = activation self . use_batch_norm = use_batch_norm if use_batch_norm : self . batch_norm_layers = nn . ModuleList ( [ nn . BatchNorm1d ( features , eps = 1e-3 ) for _ in range ( 2 )] ) if context_features is not None : self . context_layer = nn . Linear ( context_features , features ) self . linear_layers = nn . ModuleList ( [ nn . Linear ( features , features ) for _ in range ( 2 )] ) self . dropout = nn . Dropout ( p = dropout_probability ) if zero_initialization : init . uniform_ ( self . linear_layers [ - 1 ] . weight , - 1e-3 , 1e-3 ) init . uniform_ ( self . linear_layers [ - 1 ] . bias , - 1e-3 , 1e-3 ) def forward ( self , inputs , context = None ): temps = inputs if self . use_batch_norm : temps = self . batch_norm_layers [ 0 ]( temps ) temps = self . activation ( temps ) temps = self . linear_layers [ 0 ]( temps ) if self . use_batch_norm : temps = self . batch_norm_layers [ 1 ]( temps ) temps = self . activation ( temps ) temps = self . dropout ( temps ) temps = self . linear_layers [ 1 ]( temps ) if context is not None : temps = F . glu ( torch . cat (( temps , self . context_layer ( context )), dim = 1 ), dim = 1 ) return inputs + temps ResidualNet Bases: nn . Module A general-purpose residual network. Works only with 1-dim inputs. Source code in normflows/nets/resnet.py 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 class ResidualNet ( nn . Module ): \"\"\"A general-purpose residual network. Works only with 1-dim inputs.\"\"\" def __init__ ( self , in_features , out_features , hidden_features , context_features = None , num_blocks = 2 , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , preprocessing = None , ): super () . __init__ () self . hidden_features = hidden_features self . context_features = context_features self . preprocessing = preprocessing if context_features is not None : self . initial_layer = nn . Linear ( in_features + context_features , hidden_features ) else : self . initial_layer = nn . Linear ( in_features , hidden_features ) self . blocks = nn . ModuleList ( [ ResidualBlock ( features = hidden_features , context_features = context_features , activation = activation , dropout_probability = dropout_probability , use_batch_norm = use_batch_norm , ) for _ in range ( num_blocks ) ] ) self . final_layer = nn . Linear ( hidden_features , out_features ) def forward ( self , inputs , context = None ): if self . preprocessing is None : temps = inputs else : temps = self . preprocessing ( inputs ) if context is None : temps = self . initial_layer ( temps ) else : temps = self . initial_layer ( torch . cat (( temps , context ), dim = 1 )) for block in self . blocks : temps = block ( temps , context = context ) outputs = self . final_layer ( temps ) return outputs sampling hais HAIS Class which performs HAIS Source code in normflows/sampling/hais.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class HAIS : \"\"\" Class which performs HAIS \"\"\" def __init__ ( self , betas , prior , target , num_leapfrog , step_size , log_mass ): \"\"\" Args: betas: Annealing schedule, the jth target is ```f_j(x) = f_0(x)^{\\beta_j} f_n(x)^{1-\\beta_j}``` where the target is proportional to f_0 and the prior is proportional to f_n. The number of intermediate steps is infered from the shape of betas. Should be of the form 1 = \\beta_0 > \\beta_1 > ... > \\beta_n = 0 prior: The prior distribution to start the HAIS chain. target: The target distribution from which we would like to draw weighted samples. num_leapfrog: Number of leapfrog steps in the HMC transitions. step_size: step_size to use for HMC transitions. log_mass: log_mass to use for HMC transitions. \"\"\" self . prior = prior self . target = target self . layers = [] n = betas . shape [ 0 ] - 1 for i in range ( n - 1 , 0 , - 1 ): intermediate_target = distributions . LinearInterpolation ( self . target , self . prior , betas [ i ] ) self . layers += [ flows . HamiltonianMonteCarlo ( intermediate_target , num_leapfrog , torch . log ( step_size ), log_mass ) ] def sample ( self , num_samples ): \"\"\"Run HAIS to draw samples from the target with appropriate weights. Args: num_samples: The number of samples to draw.a \"\"\" samples , log_weights = self . prior . forward ( num_samples ) log_weights = - log_weights for i in range ( len ( self . layers )): samples , log_weights_addition = self . layers [ i ] . forward ( samples ) log_weights += log_weights_addition log_weights += self . target . log_prob ( samples ) return samples , log_weights __init__ ( betas , prior , target , num_leapfrog , step_size , log_mass ) Parameters: Name Type Description Default betas Annealing schedule, the jth target is f_j(x) = f_0(x)^{\beta_j} f_n(x)^{1-\beta_j} where the target is proportional to f_0 and the prior is proportional to f_n. The number of intermediate steps is infered from the shape of betas. Should be of the form 1 = \beta_0 > \beta_1 > ... > \beta_n = 0 required prior The prior distribution to start the HAIS chain. required target The target distribution from which we would like to draw weighted samples. required num_leapfrog Number of leapfrog steps in the HMC transitions. required step_size step_size to use for HMC transitions. required log_mass log_mass to use for HMC transitions. required Source code in normflows/sampling/hais.py 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 def __init__ ( self , betas , prior , target , num_leapfrog , step_size , log_mass ): \"\"\" Args: betas: Annealing schedule, the jth target is ```f_j(x) = f_0(x)^{\\beta_j} f_n(x)^{1-\\beta_j}``` where the target is proportional to f_0 and the prior is proportional to f_n. The number of intermediate steps is infered from the shape of betas. Should be of the form 1 = \\beta_0 > \\beta_1 > ... > \\beta_n = 0 prior: The prior distribution to start the HAIS chain. target: The target distribution from which we would like to draw weighted samples. num_leapfrog: Number of leapfrog steps in the HMC transitions. step_size: step_size to use for HMC transitions. log_mass: log_mass to use for HMC transitions. \"\"\" self . prior = prior self . target = target self . layers = [] n = betas . shape [ 0 ] - 1 for i in range ( n - 1 , 0 , - 1 ): intermediate_target = distributions . LinearInterpolation ( self . target , self . prior , betas [ i ] ) self . layers += [ flows . HamiltonianMonteCarlo ( intermediate_target , num_leapfrog , torch . log ( step_size ), log_mass ) ] sample ( num_samples ) Run HAIS to draw samples from the target with appropriate weights. Parameters: Name Type Description Default num_samples The number of samples to draw.a required Source code in normflows/sampling/hais.py 37 38 39 40 41 42 43 44 45 46 47 48 49 def sample ( self , num_samples ): \"\"\"Run HAIS to draw samples from the target with appropriate weights. Args: num_samples: The number of samples to draw.a \"\"\" samples , log_weights = self . prior . forward ( num_samples ) log_weights = - log_weights for i in range ( len ( self . layers )): samples , log_weights_addition = self . layers [ i ] . forward ( samples ) log_weights += log_weights_addition log_weights += self . target . log_prob ( samples ) return samples , log_weights transforms Logit Bases: flows . Flow Logit mapping of image tensor, see RealNVP paper logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) Source code in normflows/transforms.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class Logit ( flows . Flow ): \"\"\"Logit mapping of image tensor, see RealNVP paper ``` logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) ``` \"\"\" def __init__ ( self , alpha = 0.05 ): \"\"\"Constructor Args: alpha: Alpha parameter, see above \"\"\" super () . __init__ () self . alpha = alpha def forward ( self , z ): beta = 1 - 2 * self . alpha sum_dims = list ( range ( 1 , z . dim ())) ls = torch . sum ( torch . nn . functional . logsigmoid ( z ), dim = sum_dims ) mls = torch . sum ( torch . nn . functional . logsigmoid ( - z ), dim = sum_dims ) log_det = - np . log ( beta ) * np . prod ([ * z . shape [ 1 :]]) + ls + mls z = ( torch . sigmoid ( z ) - self . alpha ) / beta return z , log_det def inverse ( self , z ): beta = 1 - 2 * self . alpha z = self . alpha + beta * z logz = torch . log ( z ) log1mz = torch . log ( 1 - z ) z = logz - log1mz sum_dims = list ( range ( 1 , z . dim ())) log_det = ( np . log ( beta ) * np . prod ([ * z . shape [ 1 :]]) - torch . sum ( logz , dim = sum_dims ) - torch . sum ( log1mz , dim = sum_dims ) ) return z , log_det __init__ ( alpha = 0.05 ) Constructor Parameters: Name Type Description Default alpha Alpha parameter, see above 0.05 Source code in normflows/transforms.py 17 18 19 20 21 22 23 24 def __init__ ( self , alpha = 0.05 ): \"\"\"Constructor Args: alpha: Alpha parameter, see above \"\"\" super () . __init__ () self . alpha = alpha Shift Bases: flows . Flow Shift data by a fixed constant Default is -0.5 to shift data from interval [0, 1] to [-0.5, 0.5] Source code in normflows/transforms.py 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 class Shift ( flows . Flow ): \"\"\"Shift data by a fixed constant Default is -0.5 to shift data from interval [0, 1] to [-0.5, 0.5] \"\"\" def __init__ ( self , shift =- 0.5 ): \"\"\"Constructor Args: shift: Shift to apply to the data \"\"\" super () . __init__ () self . shift = shift def forward ( self , z ): z -= self . shift log_det = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) return z , log_det def inverse ( self , z ): z += self . shift log_det = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) return z , log_det __init__ ( shift =- 0.5 ) Constructor Parameters: Name Type Description Default shift Shift to apply to the data -0.5 Source code in normflows/transforms.py 57 58 59 60 61 62 63 64 def __init__ ( self , shift =- 0.5 ): \"\"\"Constructor Args: shift: Shift to apply to the data \"\"\" super () . __init__ () self . shift = shift transforms_test utils eval bitsPerDim ( model , x , y = None , trans = 'logit' , trans_param = [ 0.05 ]) Computes the bits per dim for a batch of data Parameters: Name Type Description Default model Model to compute bits per dim for required x Batch of data required y Class labels for batch of data if base distribution is class conditional None trans Transformation to be applied to images during training 'logit' trans_param List of parameters of the transformation [0.05] Returns: Type Description Bits per dim for data batch under model Source code in normflows/utils/eval.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 def bitsPerDim ( model , x , y = None , trans = \"logit\" , trans_param = [ 0.05 ]): \"\"\"Computes the bits per dim for a batch of data Args: model: Model to compute bits per dim for x: Batch of data y: Class labels for batch of data if base distribution is class conditional trans: Transformation to be applied to images during training trans_param: List of parameters of the transformation Returns: Bits per dim for data batch under model \"\"\" dims = torch . prod ( torch . tensor ( x . size ()[ 1 :])) if trans == \"logit\" : if y is None : log_q = model . log_prob ( x ) else : log_q = model . log_prob ( x , y ) sum_dims = list ( range ( 1 , x . dim ())) ls = torch . nn . LogSigmoid () sig_ = torch . sum ( ls ( x ) / np . log ( 2 ), sum_dims ) sig_ += torch . sum ( ls ( - x ) / np . log ( 2 ), sum_dims ) b = - log_q / dims / np . log ( 2 ) - np . log2 ( 1 - trans_param [ 0 ]) + 8 b += sig_ / dims else : raise NotImplementedError ( \"The transformation \" + trans + \" is not implemented.\" ) return b bitsPerDimDataset ( model , data_loader , class_cond = True , trans = 'logit' , trans_param = [ 0.05 ]) Computes average bits per dim for an entire dataset given by a data loader Parameters: Name Type Description Default model Model to compute bits per dim for required data_loader Data loader of dataset required class_cond Flag indicating whether model is class_conditional True trans Transformation to be applied to images during training 'logit' trans_param List of parameters of the transformation [0.05] Returns: Type Description Average bits per dim for dataset Source code in normflows/utils/eval.py 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 def bitsPerDimDataset ( model , data_loader , class_cond = True , trans = \"logit\" , trans_param = [ 0.05 ] ): \"\"\"Computes average bits per dim for an entire dataset given by a data loader Args: model: Model to compute bits per dim for data_loader: Data loader of dataset class_cond: Flag indicating whether model is class_conditional trans: Transformation to be applied to images during training trans_param: List of parameters of the transformation Returns: Average bits per dim for dataset \"\"\" n = 0 b_cum = 0 with torch . no_grad (): for x , y in iter ( data_loader ): b_ = bitsPerDim ( model , x , y . to ( x . device ) if class_cond else None , trans , trans_param ) b_np = b_ . to ( \"cpu\" ) . numpy () b_cum += np . nansum ( b_np ) n += len ( x ) - np . sum ( np . isnan ( b_np )) b = b_cum / n return b masks create_alternating_binary_mask ( features , even = True ) Creates a binary mask of a given dimension which alternates its masking. Parameters: Name Type Description Default features Dimension of mask. required even If True, even values are assigned 1s, odd 0s. If False, vice versa. True Returns: Type Description Alternating binary mask of type torch.Tensor. Source code in normflows/utils/masks.py 4 5 6 7 8 9 10 11 12 13 14 15 16 17 def create_alternating_binary_mask ( features , even = True ): \"\"\"Creates a binary mask of a given dimension which alternates its masking. Args: features: Dimension of mask. even: If True, even values are assigned 1s, odd 0s. If False, vice versa. Returns: Alternating binary mask of type torch.Tensor. \"\"\" mask = torch . zeros ( features ) . byte () start = 0 if even else 1 mask [ start :: 2 ] += 1 return mask create_mid_split_binary_mask ( features ) Creates a binary mask of a given dimension which splits its masking at the midpoint. Parameters: Name Type Description Default features Dimension of mask. required Returns: Type Description Binary mask split at midpoint of type torch.Tensor Source code in normflows/utils/masks.py 20 21 22 23 24 25 26 27 28 29 30 31 32 def create_mid_split_binary_mask ( features ): \"\"\"Creates a binary mask of a given dimension which splits its masking at the midpoint. Args: features: Dimension of mask. Returns: Binary mask split at midpoint of type torch.Tensor \"\"\" mask = torch . zeros ( features ) . byte () midpoint = features // 2 if features % 2 == 0 else features // 2 + 1 mask [: midpoint ] += 1 return mask create_random_binary_mask ( features , seed = None ) Creates a random binary mask of a given dimension with half of its entries randomly set to 1s. Parameters: Name Type Description Default features Dimension of mask. required seed Seed to be used None Returns: Type Description Binary mask with half of its entries set to 1s, of type torch.Tensor. Source code in normflows/utils/masks.py 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 def create_random_binary_mask ( features , seed = None ): \"\"\"Creates a random binary mask of a given dimension with half of its entries randomly set to 1s. Args: features: Dimension of mask. seed: Seed to be used Returns: Binary mask with half of its entries set to 1s, of type torch.Tensor. \"\"\" mask = torch . zeros ( features ) . byte () weights = torch . ones ( features ) . float () num_samples = features // 2 if features % 2 == 0 else features // 2 + 1 if seed is None : generator = None else : generator = torch . Generator () generator . manual_seed ( seed ) indices = torch . multinomial ( input = weights , num_samples = num_samples , replacement = False , generator = generator ) mask [ indices ] += 1 return mask nn ActNorm Bases: nn . Module ActNorm layer with just one forward pass Source code in normflows/utils/nn.py 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 class ActNorm ( nn . Module ): \"\"\" ActNorm layer with just one forward pass \"\"\" def __init__ ( self , shape ): \"\"\"Constructor Args: shape: Same as shape in flows.ActNorm logscale_factor: Same as shape in flows.ActNorm \"\"\" super () . __init__ () self . actNorm = flows . ActNorm ( shape ) def forward ( self , input ): out , _ = self . actNorm ( input ) return out __init__ ( shape ) Constructor Parameters: Name Type Description Default shape Same as shape in flows.ActNorm required logscale_factor Same as shape in flows.ActNorm required Source code in normflows/utils/nn.py 30 31 32 33 34 35 36 37 38 39 def __init__ ( self , shape ): \"\"\"Constructor Args: shape: Same as shape in flows.ActNorm logscale_factor: Same as shape in flows.ActNorm \"\"\" super () . __init__ () self . actNorm = flows . ActNorm ( shape ) ClampExp Bases: nn . Module Nonlinearity min(exp(lam * x), 1) Source code in normflows/utils/nn.py 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 class ClampExp ( nn . Module ): \"\"\" Nonlinearity min(exp(lam * x), 1) \"\"\" def __init__ ( self ): \"\"\"Constructor Args: lam: Lambda parameter \"\"\" super ( ClampExp , self ) . __init__ () def forward ( self , x ): one = torch . tensor ( 1.0 , device = x . device , dtype = x . dtype ) return torch . min ( torch . exp ( x ), one ) __init__ () Constructor Parameters: Name Type Description Default lam Lambda parameter required Source code in normflows/utils/nn.py 51 52 53 54 55 56 57 def __init__ ( self ): \"\"\"Constructor Args: lam: Lambda parameter \"\"\" super ( ClampExp , self ) . __init__ () ConstScaleLayer Bases: nn . Module Scaling features by a fixed factor Source code in normflows/utils/nn.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 class ConstScaleLayer ( nn . Module ): \"\"\" Scaling features by a fixed factor \"\"\" def __init__ ( self , scale = 1.0 ): \"\"\"Constructor Args: scale: Scale to apply to features \"\"\" super () . __init__ () self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu ) def forward ( self , input ): return input * self . scale __init__ ( scale = 1.0 ) Constructor Parameters: Name Type Description Default scale Scale to apply to features 1.0 Source code in normflows/utils/nn.py 12 13 14 15 16 17 18 19 20 def __init__ ( self , scale = 1.0 ): \"\"\"Constructor Args: scale: Scale to apply to features \"\"\" super () . __init__ () self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu ) PeriodicFeaturesCat Bases: nn . Module Converts a specified part of the input to periodic features by replacing those features f with [sin(scale * f), cos(scale * f)]. Note that this decreases the number of features and their order is changed. Source code in normflows/utils/nn.py 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 class PeriodicFeaturesCat ( nn . Module ): \"\"\" Converts a specified part of the input to periodic features by replacing those features f with [sin(scale * f), cos(scale * f)]. Note that this decreases the number of features and their order is changed. \"\"\" def __init__ ( self , ndim , ind , scale = 1.0 ): \"\"\" Constructor :param ndim: Int, number of dimensions :param ind: Iterable, indices of input elements to convert to periodic features :param scale: Scalar or iterable, used to scale inputs before converting them to periodic features \"\"\" super ( PeriodicFeaturesCat , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale def forward ( self , inputs ): inputs_ = inputs [ ... , self . ind ] inputs_ = self . scale * inputs_ inputs_sin = torch . sin ( inputs_ ) inputs_cos = torch . cos ( inputs_ ) out = torch . cat (( inputs_sin , inputs_cos , inputs [ ... , self . ind_ ]), - 1 ) return out __init__ ( ndim , ind , scale = 1.0 ) Constructor :param ndim: Int, number of dimensions :param ind: Iterable, indices of input elements to convert to periodic features :param scale: Scalar or iterable, used to scale inputs before converting them to periodic features Source code in normflows/utils/nn.py 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 def __init__ ( self , ndim , ind , scale = 1.0 ): \"\"\" Constructor :param ndim: Int, number of dimensions :param ind: Iterable, indices of input elements to convert to periodic features :param scale: Scalar or iterable, used to scale inputs before converting them to periodic features \"\"\" super ( PeriodicFeaturesCat , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale PeriodicFeaturesElementwise Bases: nn . Module Converts a specified part of the input to periodic features by replacing those features f with w1 * sin(scale * f) + w2 * cos(scale * f). Note that this operation is done elementwise and, therefore, some information about the feature can be lost. Source code in normflows/utils/nn.py 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 class PeriodicFeaturesElementwise ( nn . Module ): \"\"\" Converts a specified part of the input to periodic features by replacing those features f with w1 * sin(scale * f) + w2 * cos(scale * f). Note that this operation is done elementwise and, therefore, some information about the feature can be lost. \"\"\" def __init__ ( self , ndim , ind , scale = 1.0 , bias = False , activation = None ): \"\"\"Constructor Args: ndim (int): number of dimensions ind (iterable): indices of input elements to convert to periodic features scale: Scalar or iterable, used to scale inputs before converting them to periodic features bias: Flag, whether to add a bias activation: Function or None, activation function to be applied \"\"\" super ( PeriodicFeaturesElementwise , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) self . weights = nn . Parameter ( torch . ones ( len ( self . ind ), 2 )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale self . apply_bias = bias if self . apply_bias : self . bias = nn . Parameter ( torch . zeros ( len ( self . ind ))) if activation is None : self . activation = lambda input : input else : self . activation = activation def forward ( self , inputs ): inputs_ = inputs [ ... , self . ind ] inputs_ = self . scale * inputs_ inputs_ = self . weights [:, 0 ] * torch . sin ( inputs_ ) + self . weights [ :, 1 ] * torch . cos ( inputs_ ) if self . apply_bias : inputs_ = inputs_ + self . bias inputs_ = self . activation ( inputs_ ) out = torch . cat (( inputs_ , inputs [ ... , self . ind_ ]), - 1 ) return out [ ... , self . inv_perm ] __init__ ( ndim , ind , scale = 1.0 , bias = False , activation = None ) Constructor Parameters: Name Type Description Default ndim int number of dimensions required ind iterable indices of input elements to convert to periodic features required scale Scalar or iterable, used to scale inputs before converting them to periodic features 1.0 bias Flag, whether to add a bias False activation Function or None, activation function to be applied None Source code in normflows/utils/nn.py 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 def __init__ ( self , ndim , ind , scale = 1.0 , bias = False , activation = None ): \"\"\"Constructor Args: ndim (int): number of dimensions ind (iterable): indices of input elements to convert to periodic features scale: Scalar or iterable, used to scale inputs before converting them to periodic features bias: Flag, whether to add a bias activation: Function or None, activation function to be applied \"\"\" super ( PeriodicFeaturesElementwise , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) self . weights = nn . Parameter ( torch . ones ( len ( self . ind ), 2 )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale self . apply_bias = bias if self . apply_bias : self . bias = nn . Parameter ( torch . zeros ( len ( self . ind ))) if activation is None : self . activation = lambda input : input else : self . activation = activation sum_except_batch ( x , num_batch_dims = 1 ) Sums all elements of x except for the first num_batch_dims dimensions. Source code in normflows/utils/nn.py 190 191 192 193 def sum_except_batch ( x , num_batch_dims = 1 ): \"\"\"Sums all elements of `x` except for the first `num_batch_dims` dimensions.\"\"\" reduce_dims = list ( range ( num_batch_dims , x . ndimension ())) return torch . sum ( x , dim = reduce_dims ) optim clear_grad ( model ) Set gradients of model parameter to None as this speeds up training, See youtube Parameters: Name Type Description Default model Model to clear gradients of required Source code in normflows/utils/optim.py 16 17 18 19 20 21 22 23 24 25 def clear_grad ( model ): \"\"\"Set gradients of model parameter to None as this speeds up training, See [youtube](https://www.youtube.com/watch?v=9mS1fIYj1So) Args: model: Model to clear gradients of \"\"\" for param in model . parameters (): param . grad = None set_requires_grad ( module , flag ) Sets requires_grad flag of all parameters of a torch.nn.module Parameters: Name Type Description Default module torch.nn.module required flag Flag to set requires_grad to required Source code in normflows/utils/optim.py 4 5 6 7 8 9 10 11 12 13 def set_requires_grad ( module , flag ): \"\"\"Sets requires_grad flag of all parameters of a torch.nn.module Args: module: torch.nn.module flag: Flag to set requires_grad to \"\"\" for param in module . parameters (): param . requires_grad = flag preprocessing Jitter Transform for dataloader, adds uniform jitter noise to data Source code in normflows/utils/preprocessing.py 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 class Jitter : \"\"\"Transform for dataloader, adds uniform jitter noise to data\"\"\" def __init__ ( self , scale = 1.0 / 256 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale def __call__ ( self , x ): eps = torch . rand_like ( x ) * self . scale x_ = x + eps return x_ __init__ ( scale = 1.0 / 256 ) Constructor Parameters: Name Type Description Default scale Scaling factor for noise 1.0 / 256 Source code in normflows/utils/preprocessing.py 31 32 33 34 35 36 37 def __init__ ( self , scale = 1.0 / 256 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale Logit Transform for dataloader logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) Source code in normflows/utils/preprocessing.py 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 class Logit : \"\"\"Transform for dataloader ``` logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) ``` \"\"\" def __init__ ( self , alpha = 0 ): \"\"\"Constructor Args: alpha: see above \"\"\" self . alpha = alpha def __call__ ( self , x ): x_ = self . alpha + ( 1 - self . alpha ) * x return torch . log ( x_ / ( 1 - x_ )) def inverse ( self , x ): return ( torch . sigmoid ( x ) - self . alpha ) / ( 1 - self . alpha ) __init__ ( alpha = 0 ) Constructor Parameters: Name Type Description Default alpha see above 0 Source code in normflows/utils/preprocessing.py 12 13 14 15 16 17 18 def __init__ ( self , alpha = 0 ): \"\"\"Constructor Args: alpha: see above \"\"\" self . alpha = alpha Scale Transform for dataloader, adds uniform jitter noise to data Source code in normflows/utils/preprocessing.py 45 46 47 48 49 50 51 52 53 54 55 56 57 class Scale : \"\"\"Transform for dataloader, adds uniform jitter noise to data\"\"\" def __init__ ( self , scale = 255.0 / 256.0 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale def __call__ ( self , x ): return x * self . scale __init__ ( scale = 255.0 / 256.0 ) Constructor Parameters: Name Type Description Default scale Scaling factor for noise 255.0 / 256.0 Source code in normflows/utils/preprocessing.py 48 49 50 51 52 53 54 def __init__ ( self , scale = 255.0 / 256.0 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale","title":"API"},{"location":"references/#api-references","text":"","title":"API references"},{"location":"references/#normflows.core","text":"","title":"core"},{"location":"references/#normflows.core.ClassCondFlow","text":"Bases: nn . Module Class conditional normalizing Flow model Source code in normflows/core.py 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 class ClassCondFlow ( nn . Module ): \"\"\" Class conditional normalizing Flow model \"\"\" def __init__ ( self , q0 , flows ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) def forward_kld ( self , x , y ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return - torch . mean ( log_q ) def sample ( self , num_samples = 1 , y = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples , y ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return log_q def save ( self , path ): \"\"\"Save state dict of model Args: param path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path ))","title":"ClassCondFlow"},{"location":"references/#normflows.core.ClassCondFlow.__init__","text":"Constructor Parameters: Name Type Description Default q0 Base distribution required flows List of flows required Source code in normflows/core.py 221 222 223 224 225 226 227 228 229 230 def __init__ ( self , q0 , flows ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows )","title":"__init__()"},{"location":"references/#normflows.core.ClassCondFlow.forward_kld","text":"Estimates forward KL divergence, see arXiv 1912.02762 Parameters: Name Type Description Default x Batch sampled from target distribution required Returns: Type Description Estimate of forward KL divergence averaged over batch Source code in normflows/core.py 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 def forward_kld ( self , x , y ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return - torch . mean ( log_q )","title":"forward_kld()"},{"location":"references/#normflows.core.ClassCondFlow.load","text":"Load model from state dict Parameters: Name Type Description Default path Path including filename where to load model from required Source code in normflows/core.py 291 292 293 294 295 296 297 def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path ))","title":"load()"},{"location":"references/#normflows.core.ClassCondFlow.log_prob","text":"Get log probability for batch Parameters: Name Type Description Default x Batch required y Classes of x required Returns: Type Description log probability Source code in normflows/core.py 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z , y ) return log_q","title":"log_prob()"},{"location":"references/#normflows.core.ClassCondFlow.sample","text":"Samples from flow-based approximate distribution Parameters: Name Type Description Default num_samples Number of samples to draw 1 y Classes to sample from, will be sampled uniformly if None None Returns: Type Description Samples, log probability Source code in normflows/core.py 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 def sample ( self , num_samples = 1 , y = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples , y ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q","title":"sample()"},{"location":"references/#normflows.core.ClassCondFlow.save","text":"Save state dict of model Parameters: Name Type Description Default param path Path including filename where to save model required Source code in normflows/core.py 283 284 285 286 287 288 289 def save ( self , path ): \"\"\"Save state dict of model Args: param path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path )","title":"save()"},{"location":"references/#normflows.core.MultiscaleFlow","text":"Bases: nn . Module Normalizing Flow model with multiscale architecture, see RealNVP or Glow paper Source code in normflows/core.py 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 class MultiscaleFlow ( nn . Module ): \"\"\" Normalizing Flow model with multiscale architecture, see RealNVP or Glow paper \"\"\" def __init__ ( self , q0 , flows , merges , transform = None , class_cond = True ): \"\"\"Constructor Args: q0: List of base distribution flows: List of list of flows for each level merges: List of merge/split operations (forward pass must do merge) transform: Initial transformation of inputs class_cond: Flag, indicated whether model has class conditional base distributions \"\"\" super () . __init__ () self . q0 = nn . ModuleList ( q0 ) self . num_levels = len ( self . q0 ) self . flows = torch . nn . ModuleList ([ nn . ModuleList ( flow ) for flow in flows ]) self . merges = torch . nn . ModuleList ( merges ) self . transform = transform self . class_cond = class_cond def forward_kld ( self , x , y = None ): \"\"\"Estimates forward KL divergence, see see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution y: Batch of targets, if applicable Returns: Estimate of forward KL divergence averaged over batch \"\"\" return - torch . mean ( self . log_prob ( x , y )) def forward ( self , x , y = None ): \"\"\"Get negative log-likelihood for maximum likelihood training Args: x: Batch of data y: Batch of targets, if applicable Returns: Negative log-likelihood of the batch \"\"\" return - self . log_prob ( x , y ) def sample ( self , num_samples = 1 , y = None , temperature = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None temperature: Temperature parameter for temp annealed sampling Returns: Samples, log probability \"\"\" if temperature is not None : self . set_temperature ( temperature ) for i in range ( len ( self . q0 )): if self . class_cond : z_ , log_q_ = self . q0 [ i ]( num_samples , y ) else : z_ , log_q_ = self . q0 [ i ]( num_samples ) if i == 0 : log_q = log_q_ z = z_ else : log_q += log_q_ z , log_det = self . merges [ i - 1 ]([ z , z_ ]) log_q -= log_det for flow in self . flows [ i ]: z , log_det = flow ( z ) log_q -= log_det if self . transform is not None : z , log_det = self . transform ( z ) log_q -= log_det if temperature is not None : self . reset_temperature () return z , log_q def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = 0 z = x if self . transform is not None : z , log_det = self . transform . inverse ( z ) log_q += log_det for i in range ( len ( self . q0 ) - 1 , - 1 , - 1 ): for j in range ( len ( self . flows [ i ]) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ][ j ] . inverse ( z ) log_q += log_det if i > 0 : [ z , z_ ], log_det = self . merges [ i - 1 ] . inverse ( z ) log_q += log_det else : z_ = z if self . class_cond : log_q += self . q0 [ i ] . log_prob ( z_ , y ) else : log_q += self . q0 [ i ] . log_prob ( z_ ) return log_q def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path )) def set_temperature ( self , temperature ): \"\"\"Set temperature for temperature a annealed sampling Args: temperature: Temperature parameter \"\"\" for q0 in self . q0 : if hasattr ( q0 , \"temperature\" ): q0 . temperature = temperature else : raise NotImplementedError ( \"One base function does not \" \"support temperature annealed sampling\" ) def reset_temperature ( self ): \"\"\" Set temperature values of base distributions back to None \"\"\" self . set_temperature ( None )","title":"MultiscaleFlow"},{"location":"references/#normflows.core.MultiscaleFlow.__init__","text":"Constructor Parameters: Name Type Description Default q0 List of base distribution required flows List of list of flows for each level required merges List of merge/split operations (forward pass must do merge) required transform Initial transformation of inputs None class_cond Flag, indicated whether model has class conditional True base distributions Source code in normflows/core.py 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 def __init__ ( self , q0 , flows , merges , transform = None , class_cond = True ): \"\"\"Constructor Args: q0: List of base distribution flows: List of list of flows for each level merges: List of merge/split operations (forward pass must do merge) transform: Initial transformation of inputs class_cond: Flag, indicated whether model has class conditional base distributions \"\"\" super () . __init__ () self . q0 = nn . ModuleList ( q0 ) self . num_levels = len ( self . q0 ) self . flows = torch . nn . ModuleList ([ nn . ModuleList ( flow ) for flow in flows ]) self . merges = torch . nn . ModuleList ( merges ) self . transform = transform self . class_cond = class_cond","title":"__init__()"},{"location":"references/#normflows.core.MultiscaleFlow.forward","text":"Get negative log-likelihood for maximum likelihood training Parameters: Name Type Description Default x Batch of data required y Batch of targets, if applicable None Returns: Type Description Negative log-likelihood of the batch Source code in normflows/core.py 337 338 339 340 341 342 343 344 345 346 347 def forward ( self , x , y = None ): \"\"\"Get negative log-likelihood for maximum likelihood training Args: x: Batch of data y: Batch of targets, if applicable Returns: Negative log-likelihood of the batch \"\"\" return - self . log_prob ( x , y )","title":"forward()"},{"location":"references/#normflows.core.MultiscaleFlow.forward_kld","text":"Estimates forward KL divergence, see see arXiv 1912.02762 Parameters: Name Type Description Default x Batch sampled from target distribution required y Batch of targets, if applicable None Returns: Type Description Estimate of forward KL divergence averaged over batch Source code in normflows/core.py 325 326 327 328 329 330 331 332 333 334 335 def forward_kld ( self , x , y = None ): \"\"\"Estimates forward KL divergence, see see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution y: Batch of targets, if applicable Returns: Estimate of forward KL divergence averaged over batch \"\"\" return - torch . mean ( self . log_prob ( x , y ))","title":"forward_kld()"},{"location":"references/#normflows.core.MultiscaleFlow.load","text":"Load model from state dict Parameters: Name Type Description Default path Path including filename where to load model from required Source code in normflows/core.py 422 423 424 425 426 427 428 def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path ))","title":"load()"},{"location":"references/#normflows.core.MultiscaleFlow.log_prob","text":"Get log probability for batch Parameters: Name Type Description Default x Batch required y Classes of x required Returns: Type Description log probability Source code in normflows/core.py 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 def log_prob ( self , x , y ): \"\"\"Get log probability for batch Args: x: Batch y: Classes of x Returns: log probability \"\"\" log_q = 0 z = x if self . transform is not None : z , log_det = self . transform . inverse ( z ) log_q += log_det for i in range ( len ( self . q0 ) - 1 , - 1 , - 1 ): for j in range ( len ( self . flows [ i ]) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ][ j ] . inverse ( z ) log_q += log_det if i > 0 : [ z , z_ ], log_det = self . merges [ i - 1 ] . inverse ( z ) log_q += log_det else : z_ = z if self . class_cond : log_q += self . q0 [ i ] . log_prob ( z_ , y ) else : log_q += self . q0 [ i ] . log_prob ( z_ ) return log_q","title":"log_prob()"},{"location":"references/#normflows.core.MultiscaleFlow.reset_temperature","text":"Set temperature values of base distributions back to None Source code in normflows/core.py 445 446 447 448 449 def reset_temperature ( self ): \"\"\" Set temperature values of base distributions back to None \"\"\" self . set_temperature ( None )","title":"reset_temperature()"},{"location":"references/#normflows.core.MultiscaleFlow.sample","text":"Samples from flow-based approximate distribution Parameters: Name Type Description Default num_samples Number of samples to draw 1 y Classes to sample from, will be sampled uniformly if None None temperature Temperature parameter for temp annealed sampling None Returns: Type Description Samples, log probability Source code in normflows/core.py 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 def sample ( self , num_samples = 1 , y = None , temperature = None ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw y: Classes to sample from, will be sampled uniformly if None temperature: Temperature parameter for temp annealed sampling Returns: Samples, log probability \"\"\" if temperature is not None : self . set_temperature ( temperature ) for i in range ( len ( self . q0 )): if self . class_cond : z_ , log_q_ = self . q0 [ i ]( num_samples , y ) else : z_ , log_q_ = self . q0 [ i ]( num_samples ) if i == 0 : log_q = log_q_ z = z_ else : log_q += log_q_ z , log_det = self . merges [ i - 1 ]([ z , z_ ]) log_q -= log_det for flow in self . flows [ i ]: z , log_det = flow ( z ) log_q -= log_det if self . transform is not None : z , log_det = self . transform ( z ) log_q -= log_det if temperature is not None : self . reset_temperature () return z , log_q","title":"sample()"},{"location":"references/#normflows.core.MultiscaleFlow.save","text":"Save state dict of model Parameters: Name Type Description Default path Path including filename where to save model required Source code in normflows/core.py 414 415 416 417 418 419 420 def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path )","title":"save()"},{"location":"references/#normflows.core.MultiscaleFlow.set_temperature","text":"Set temperature for temperature a annealed sampling Parameters: Name Type Description Default temperature Temperature parameter required Source code in normflows/core.py 430 431 432 433 434 435 436 437 438 439 440 441 442 443 def set_temperature ( self , temperature ): \"\"\"Set temperature for temperature a annealed sampling Args: temperature: Temperature parameter \"\"\" for q0 in self . q0 : if hasattr ( q0 , \"temperature\" ): q0 . temperature = temperature else : raise NotImplementedError ( \"One base function does not \" \"support temperature annealed sampling\" )","title":"set_temperature()"},{"location":"references/#normflows.core.NormalizingFlow","text":"Bases: nn . Module Normalizing Flow model to approximate target distribution Source code in normflows/core.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 class NormalizingFlow ( nn . Module ): \"\"\" Normalizing Flow model to approximate target distribution \"\"\" def __init__ ( self , q0 , flows , p = None ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows p: Target distribution \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) self . p = p def forward ( self , z ): \"\"\"Transforms latent variable z to the flow variable x Args: z: Batch in the latent space Returns: Batch in the space of the target distribution \"\"\" for flow in self . flows : z , _ = flow ( z ) return z def forward_and_log_det ( self , z ): \"\"\"Transforms latent variable z to the flow variable x and computes log determinant of the Jacobian Args: z: Batch in the latent space Returns: Batch in the space of the target distribution, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( z ), device = z . device ) for flow in self . flows : z , log_d = flow ( z ) log_det += log_d return z , log_det def inverse ( self , x ): \"\"\"Transforms flow variable x to the latent variable z Args: x: Batch in the space of the target distribution Returns: Batch in the latent space \"\"\" for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , _ = self . flows [ i ] . inverse ( x ) return x def inverse_and_log_det ( self , x ): \"\"\"Transforms flow variable x to the latent variable z and computes log determinant of the Jacobian Args: x: Batch in the space of the target distribution Returns: Batch in the latent space, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( x ), device = x . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , log_d = self . flows [ i ] . inverse ( x ) log_det += log_d return x , log_det def forward_kld ( self , x ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return - torch . mean ( log_q ) def reverse_kld ( self , num_samples = 1 , beta = 1.0 , score_fn = True ): \"\"\"Estimates reverse KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: num_samples: Number of samples to draw from base distribution beta: Annealing parameter, see [arXiv 1505.05770](https://arxiv.org/abs/1505.05770) score_fn: Flag whether to include score function in gradient, see [arXiv 1703.09194](https://arxiv.org/abs/1703.09194) Returns: Estimate of the reverse KL divergence averaged over latent samples \"\"\" z , log_q_ = self . q0 ( num_samples ) log_q = torch . zeros_like ( log_q_ ) log_q += log_q_ for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det if not score_fn : z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) log_p = self . p . log_prob ( z ) return torch . mean ( log_q ) - beta * torch . mean ( log_p ) def reverse_alpha_div ( self , num_samples = 1 , alpha = 1 , dreg = False ): \"\"\"Alpha divergence when sampling from q Args: num_samples: Number of samples to draw dreg: Flag whether to use Double Reparametrized Gradient estimator, see [arXiv 1810.04152](https://arxiv.org/abs/1810.04152) Returns: Alpha divergence \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . p . log_prob ( z ) if dreg : w_const = torch . exp ( log_p - log_q ) . detach () z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) w = torch . exp ( log_p - log_q ) w_alpha = w_const ** alpha w_alpha = w_alpha / torch . mean ( w_alpha ) weights = ( 1 - alpha ) * w_alpha + alpha * w_alpha ** 2 loss = - alpha * torch . mean ( weights * torch . log ( w )) else : loss = np . sign ( alpha - 1 ) * torch . logsumexp ( alpha * ( log_p - log_q ), 0 ) return loss def sample ( self , num_samples = 1 ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q def log_prob ( self , x ): \"\"\"Get log probability for batch Args: x: Batch Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return log_q def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path ) def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path ))","title":"NormalizingFlow"},{"location":"references/#normflows.core.NormalizingFlow.__init__","text":"Constructor Parameters: Name Type Description Default q0 Base distribution required flows List of flows required p Target distribution None Source code in normflows/core.py 14 15 16 17 18 19 20 21 22 23 24 25 def __init__ ( self , q0 , flows , p = None ): \"\"\"Constructor Args: q0: Base distribution flows: List of flows p: Target distribution \"\"\" super () . __init__ () self . q0 = q0 self . flows = nn . ModuleList ( flows ) self . p = p","title":"__init__()"},{"location":"references/#normflows.core.NormalizingFlow.forward","text":"Transforms latent variable z to the flow variable x Parameters: Name Type Description Default z Batch in the latent space required Returns: Type Description Batch in the space of the target distribution Source code in normflows/core.py 27 28 29 30 31 32 33 34 35 36 37 38 def forward ( self , z ): \"\"\"Transforms latent variable z to the flow variable x Args: z: Batch in the latent space Returns: Batch in the space of the target distribution \"\"\" for flow in self . flows : z , _ = flow ( z ) return z","title":"forward()"},{"location":"references/#normflows.core.NormalizingFlow.forward_and_log_det","text":"Transforms latent variable z to the flow variable x and computes log determinant of the Jacobian Parameters: Name Type Description Default z Batch in the latent space required Returns: Type Description Batch in the space of the target distribution, log determinant of the Jacobian Source code in normflows/core.py 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def forward_and_log_det ( self , z ): \"\"\"Transforms latent variable z to the flow variable x and computes log determinant of the Jacobian Args: z: Batch in the latent space Returns: Batch in the space of the target distribution, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( z ), device = z . device ) for flow in self . flows : z , log_d = flow ( z ) log_det += log_d return z , log_det","title":"forward_and_log_det()"},{"location":"references/#normflows.core.NormalizingFlow.forward_kld","text":"Estimates forward KL divergence, see arXiv 1912.02762 Parameters: Name Type Description Default x Batch sampled from target distribution required Returns: Type Description Estimate of forward KL divergence averaged over batch Source code in normflows/core.py 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 def forward_kld ( self , x ): \"\"\"Estimates forward KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: x: Batch sampled from target distribution Returns: Estimate of forward KL divergence averaged over batch \"\"\" log_q = torch . zeros ( len ( x ), device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return - torch . mean ( log_q )","title":"forward_kld()"},{"location":"references/#normflows.core.NormalizingFlow.inverse","text":"Transforms flow variable x to the latent variable z Parameters: Name Type Description Default x Batch in the space of the target distribution required Returns: Type Description Batch in the latent space Source code in normflows/core.py 57 58 59 60 61 62 63 64 65 66 67 68 def inverse ( self , x ): \"\"\"Transforms flow variable x to the latent variable z Args: x: Batch in the space of the target distribution Returns: Batch in the latent space \"\"\" for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , _ = self . flows [ i ] . inverse ( x ) return x","title":"inverse()"},{"location":"references/#normflows.core.NormalizingFlow.inverse_and_log_det","text":"Transforms flow variable x to the latent variable z and computes log determinant of the Jacobian Parameters: Name Type Description Default x Batch in the space of the target distribution required Returns: Type Description Batch in the latent space, log determinant of the Jacobian Source code in normflows/core.py 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 def inverse_and_log_det ( self , x ): \"\"\"Transforms flow variable x to the latent variable z and computes log determinant of the Jacobian Args: x: Batch in the space of the target distribution Returns: Batch in the latent space, log determinant of the Jacobian \"\"\" log_det = torch . zeros ( len ( x ), device = x . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): x , log_d = self . flows [ i ] . inverse ( x ) log_det += log_d return x , log_det","title":"inverse_and_log_det()"},{"location":"references/#normflows.core.NormalizingFlow.load","text":"Load model from state dict Parameters: Name Type Description Default path Path including filename where to load model from required Source code in normflows/core.py 207 208 209 210 211 212 213 def load ( self , path ): \"\"\"Load model from state dict Args: path: Path including filename where to load model from \"\"\" self . load_state_dict ( torch . load ( path ))","title":"load()"},{"location":"references/#normflows.core.NormalizingFlow.log_prob","text":"Get log probability for batch Parameters: Name Type Description Default x Batch required Returns: Type Description log probability Source code in normflows/core.py 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 def log_prob ( self , x ): \"\"\"Get log probability for batch Args: x: Batch Returns: log probability \"\"\" log_q = torch . zeros ( len ( x ), dtype = x . dtype , device = x . device ) z = x for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_q += log_det log_q += self . q0 . log_prob ( z ) return log_q","title":"log_prob()"},{"location":"references/#normflows.core.NormalizingFlow.reverse_alpha_div","text":"Alpha divergence when sampling from q Parameters: Name Type Description Default num_samples Number of samples to draw 1 dreg Flag whether to use Double Reparametrized Gradient estimator, see arXiv 1810.04152 False Returns: Type Description Alpha divergence Source code in normflows/core.py 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 def reverse_alpha_div ( self , num_samples = 1 , alpha = 1 , dreg = False ): \"\"\"Alpha divergence when sampling from q Args: num_samples: Number of samples to draw dreg: Flag whether to use Double Reparametrized Gradient estimator, see [arXiv 1810.04152](https://arxiv.org/abs/1810.04152) Returns: Alpha divergence \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . p . log_prob ( z ) if dreg : w_const = torch . exp ( log_p - log_q ) . detach () z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) w = torch . exp ( log_p - log_q ) w_alpha = w_const ** alpha w_alpha = w_alpha / torch . mean ( w_alpha ) weights = ( 1 - alpha ) * w_alpha + alpha * w_alpha ** 2 loss = - alpha * torch . mean ( weights * torch . log ( w )) else : loss = np . sign ( alpha - 1 ) * torch . logsumexp ( alpha * ( log_p - log_q ), 0 ) return loss","title":"reverse_alpha_div()"},{"location":"references/#normflows.core.NormalizingFlow.reverse_kld","text":"Estimates reverse KL divergence, see arXiv 1912.02762 Parameters: Name Type Description Default num_samples Number of samples to draw from base distribution 1 beta Annealing parameter, see arXiv 1505.05770 1.0 score_fn Flag whether to include score function in gradient, see arXiv 1703.09194 True Returns: Type Description Estimate of the reverse KL divergence averaged over latent samples Source code in normflows/core.py 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 def reverse_kld ( self , num_samples = 1 , beta = 1.0 , score_fn = True ): \"\"\"Estimates reverse KL divergence, see [arXiv 1912.02762](https://arxiv.org/abs/1912.02762) Args: num_samples: Number of samples to draw from base distribution beta: Annealing parameter, see [arXiv 1505.05770](https://arxiv.org/abs/1505.05770) score_fn: Flag whether to include score function in gradient, see [arXiv 1703.09194](https://arxiv.org/abs/1703.09194) Returns: Estimate of the reverse KL divergence averaged over latent samples \"\"\" z , log_q_ = self . q0 ( num_samples ) log_q = torch . zeros_like ( log_q_ ) log_q += log_q_ for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det if not score_fn : z_ = z log_q = torch . zeros ( len ( z_ ), device = z_ . device ) utils . set_requires_grad ( self , False ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z_ , log_det = self . flows [ i ] . inverse ( z_ ) log_q += log_det log_q += self . q0 . log_prob ( z_ ) utils . set_requires_grad ( self , True ) log_p = self . p . log_prob ( z ) return torch . mean ( log_q ) - beta * torch . mean ( log_p )","title":"reverse_kld()"},{"location":"references/#normflows.core.NormalizingFlow.sample","text":"Samples from flow-based approximate distribution Parameters: Name Type Description Default num_samples Number of samples to draw 1 Returns: Type Description Samples, log probability Source code in normflows/core.py 167 168 169 170 171 172 173 174 175 176 177 178 179 180 def sample ( self , num_samples = 1 ): \"\"\"Samples from flow-based approximate distribution Args: num_samples: Number of samples to draw Returns: Samples, log probability \"\"\" z , log_q = self . q0 ( num_samples ) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det return z , log_q","title":"sample()"},{"location":"references/#normflows.core.NormalizingFlow.save","text":"Save state dict of model Parameters: Name Type Description Default path Path including filename where to save model required Source code in normflows/core.py 199 200 201 202 203 204 205 def save ( self , path ): \"\"\"Save state dict of model Args: path: Path including filename where to save model \"\"\" torch . save ( self . state_dict (), path )","title":"save()"},{"location":"references/#normflows.core.NormalizingFlowVAE","text":"Bases: nn . Module VAE using normalizing flows to express approximate distribution Source code in normflows/core.py 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 class NormalizingFlowVAE ( nn . Module ): \"\"\" VAE using normalizing flows to express approximate distribution \"\"\" def __init__ ( self , prior , q0 = distributions . Dirac (), flows = None , decoder = None ): \"\"\"Constructor of normalizing flow model Args: prior: Prior distribution of te VAE, i.e. Gaussian decoder: Optional decoder flows: Flows to transform output of base encoder q0: Base Encoder \"\"\" super () . __init__ () self . prior = prior self . decoder = decoder self . flows = nn . ModuleList ( flows ) self . q0 = q0 def forward ( self , x , num_samples = 1 ): \"\"\"Takes data batch, samples num_samples for each data point from base distribution Args: x: data batch num_samples: number of samples to draw for each data point Returns: latent variables for each batch and sample, log_q, and log_p \"\"\" z , log_q = self . q0 ( x , num_samples = num_samples ) # Flatten batch and sample dim z = z . view ( - 1 , * z . size ()[ 2 :]) log_q = log_q . view ( - 1 , * log_q . size ()[ 2 :]) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . prior . log_prob ( z ) if self . decoder is not None : log_p += self . decoder . log_prob ( x , z ) # Separate batch and sample dimension again z = z . view ( - 1 , num_samples , * z . size ()[ 1 :]) log_q = log_q . view ( - 1 , num_samples , * log_q . size ()[ 1 :]) log_p = log_p . view ( - 1 , num_samples , * log_p . size ()[ 1 :]) return z , log_q , log_p","title":"NormalizingFlowVAE"},{"location":"references/#normflows.core.NormalizingFlowVAE.__init__","text":"Constructor of normalizing flow model Parameters: Name Type Description Default prior Prior distribution of te VAE, i.e. Gaussian required decoder Optional decoder None flows Flows to transform output of base encoder None q0 Base Encoder distributions.Dirac() Source code in normflows/core.py 457 458 459 460 461 462 463 464 465 466 467 468 469 470 def __init__ ( self , prior , q0 = distributions . Dirac (), flows = None , decoder = None ): \"\"\"Constructor of normalizing flow model Args: prior: Prior distribution of te VAE, i.e. Gaussian decoder: Optional decoder flows: Flows to transform output of base encoder q0: Base Encoder \"\"\" super () . __init__ () self . prior = prior self . decoder = decoder self . flows = nn . ModuleList ( flows ) self . q0 = q0","title":"__init__()"},{"location":"references/#normflows.core.NormalizingFlowVAE.forward","text":"Takes data batch, samples num_samples for each data point from base distribution Parameters: Name Type Description Default x data batch required num_samples number of samples to draw for each data point 1 Returns: Type Description latent variables for each batch and sample, log_q, and log_p Source code in normflows/core.py 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 def forward ( self , x , num_samples = 1 ): \"\"\"Takes data batch, samples num_samples for each data point from base distribution Args: x: data batch num_samples: number of samples to draw for each data point Returns: latent variables for each batch and sample, log_q, and log_p \"\"\" z , log_q = self . q0 ( x , num_samples = num_samples ) # Flatten batch and sample dim z = z . view ( - 1 , * z . size ()[ 2 :]) log_q = log_q . view ( - 1 , * log_q . size ()[ 2 :]) for flow in self . flows : z , log_det = flow ( z ) log_q -= log_det log_p = self . prior . log_prob ( z ) if self . decoder is not None : log_p += self . decoder . log_prob ( x , z ) # Separate batch and sample dimension again z = z . view ( - 1 , num_samples , * z . size ()[ 1 :]) log_q = log_q . view ( - 1 , num_samples , * log_q . size ()[ 1 :]) log_p = log_p . view ( - 1 , num_samples , * log_p . size ()[ 1 :]) return z , log_q , log_p","title":"forward()"},{"location":"references/#normflows.core_test","text":"","title":"core_test"},{"location":"references/#normflows.distributions","text":"","title":"distributions"},{"location":"references/#normflows.distributions.base","text":"","title":"base"},{"location":"references/#normflows.distributions.base.AffineGaussian","text":"Bases: BaseDistribution Diagonal Gaussian an affine constant transformation applied to it, can be class conditional or not Source code in normflows/distributions/base.py 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 class AffineGaussian ( BaseDistribution ): \"\"\" Diagonal Gaussian an affine constant transformation applied to it, can be class conditional or not \"\"\" def __init__ ( self , shape , affine_shape , num_classes = None ): \"\"\"Constructor Args: shape: Shape of the variables affine_shape: Shape of the parameters in the affine transformation num_classes: Number of classes if the base is class conditional, None otherwise \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . affine_shape = affine_shape self . num_classes = num_classes self . class_cond = num_classes is not None # Affine transformation if self . class_cond : self . transform = flows . CCAffineConst ( self . affine_shape , self . num_classes ) else : self . transform = flows . AffineConstFlow ( self . affine_shape ) # Temperature parameter for annealed sampling self . temperature = None def forward ( self , num_samples = 1 , y = None ): dtype = self . transform . s . dtype device = self . transform . s . device if self . class_cond : if y is not None : num_samples = len ( y ) else : y = torch . randint ( self . num_classes , ( num_samples ,), device = device ) if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = dtype , device = device ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot if self . temperature is not None : log_scale = np . log ( self . temperature ) else : log_scale = 0.0 # Sample eps = torch . randn (( num_samples ,) + self . shape , dtype = dtype , device = device ) z = np . exp ( log_scale ) * eps # Get log prob log_p = ( - 0.5 * self . d * np . log ( 2 * np . pi ) - self . d * log_scale - 0.5 * torch . sum ( torch . pow ( eps , 2 ), dim = self . sum_dim ) ) # Apply transform if self . class_cond : z , log_det = self . transform ( z , y ) else : z , log_det = self . transform ( z ) log_p -= log_det return z , log_p def log_prob ( self , z , y = None ): # Perpare onehot encoding of class if needed if self . class_cond : if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = self . transform . s . dtype , device = self . transform . s . device , ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot if self . temperature is not None : log_scale = np . log ( self . temperature ) else : log_scale = 0.0 # Get log prob if self . class_cond : z , log_p = self . transform . inverse ( z , y ) else : z , log_p = self . transform . inverse ( z ) z = z / np . exp ( log_scale ) log_p = ( log_p - self . d * log_scale - 0.5 * self . d * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . pow ( z , 2 ), dim = self . sum_dim ) ) return log_p","title":"AffineGaussian"},{"location":"references/#normflows.distributions.base.AffineGaussian.__init__","text":"Constructor Parameters: Name Type Description Default shape Shape of the variables required affine_shape Shape of the parameters in the affine transformation required num_classes Number of classes if the base is class conditional, None otherwise None Source code in normflows/distributions/base.py 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 def __init__ ( self , shape , affine_shape , num_classes = None ): \"\"\"Constructor Args: shape: Shape of the variables affine_shape: Shape of the parameters in the affine transformation num_classes: Number of classes if the base is class conditional, None otherwise \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . affine_shape = affine_shape self . num_classes = num_classes self . class_cond = num_classes is not None # Affine transformation if self . class_cond : self . transform = flows . CCAffineConst ( self . affine_shape , self . num_classes ) else : self . transform = flows . AffineConstFlow ( self . affine_shape ) # Temperature parameter for annealed sampling self . temperature = None","title":"__init__()"},{"location":"references/#normflows.distributions.base.BaseDistribution","text":"Bases: nn . Module Base distribution of a flow-based model Parameters do not depend of target variable (as is the case for a VAE encoder) Source code in normflows/distributions/base.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class BaseDistribution ( nn . Module ): \"\"\" Base distribution of a flow-based model Parameters do not depend of target variable (as is the case for a VAE encoder) \"\"\" def __init__ ( self ): super () . __init__ () def forward ( self , num_samples = 1 ): \"\"\"Samples from base distribution and calculates log probability Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution, log probability \"\"\" raise NotImplementedError def log_prob ( self , z ): \"\"\"Calculate log probability of batch of samples Args: z: Batch of random variables to determine log probability for Returns: log probability for each batch element \"\"\" raise NotImplementedError def sample ( self , num_samples = 1 , ** kwargs ): \"\"\"Samples from base distribution Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution \"\"\" z , _ = self . forward ( num_samples , ** kwargs ) return z","title":"BaseDistribution"},{"location":"references/#normflows.distributions.base.BaseDistribution.forward","text":"Samples from base distribution and calculates log probability Parameters: Name Type Description Default num_samples Number of samples to draw from the distriubtion 1 Returns: Type Description Samples drawn from the distribution, log probability Source code in normflows/distributions/base.py 17 18 19 20 21 22 23 24 25 26 def forward ( self , num_samples = 1 ): \"\"\"Samples from base distribution and calculates log probability Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution, log probability \"\"\" raise NotImplementedError","title":"forward()"},{"location":"references/#normflows.distributions.base.BaseDistribution.log_prob","text":"Calculate log probability of batch of samples Parameters: Name Type Description Default z Batch of random variables to determine log probability for required Returns: Type Description log probability for each batch element Source code in normflows/distributions/base.py 28 29 30 31 32 33 34 35 36 37 def log_prob ( self , z ): \"\"\"Calculate log probability of batch of samples Args: z: Batch of random variables to determine log probability for Returns: log probability for each batch element \"\"\" raise NotImplementedError","title":"log_prob()"},{"location":"references/#normflows.distributions.base.BaseDistribution.sample","text":"Samples from base distribution Parameters: Name Type Description Default num_samples Number of samples to draw from the distriubtion 1 Returns: Type Description Samples drawn from the distribution Source code in normflows/distributions/base.py 39 40 41 42 43 44 45 46 47 48 49 def sample ( self , num_samples = 1 , ** kwargs ): \"\"\"Samples from base distribution Args: num_samples: Number of samples to draw from the distriubtion Returns: Samples drawn from the distribution \"\"\" z , _ = self . forward ( num_samples , ** kwargs ) return z","title":"sample()"},{"location":"references/#normflows.distributions.base.ClassCondDiagGaussian","text":"Bases: BaseDistribution Class conditional multivariate Gaussian distribution with diagonal covariance matrix Source code in normflows/distributions/base.py 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 class ClassCondDiagGaussian ( BaseDistribution ): \"\"\" Class conditional multivariate Gaussian distribution with diagonal covariance matrix \"\"\" def __init__ ( self , shape , num_classes ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension num_classes: Number of classes to condition on \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . perm = [ self . n_dim ] + list ( range ( self . n_dim )) self . d = np . prod ( shape ) self . num_classes = num_classes self . loc = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . log_scale = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . temperature = None # Temperature parameter for annealed sampling def forward ( self , num_samples = 1 , y = None ): if y is not None : num_samples = len ( y ) else : y = torch . randint ( self . num_classes , ( num_samples ,), device = self . loc . device ) if y . dim () == 1 : y_onehot = torch . zeros ( ( self . num_classes , num_samples ), dtype = self . loc . dtype , device = self . loc . device , ) y_onehot . scatter_ ( 0 , y [ None ], 1 ) y = y_onehot else : y = y . t () eps = torch . randn ( ( num_samples ,) + self . shape , dtype = self . loc . dtype , device = self . loc . device ) loc = ( self . loc @ y ) . permute ( * self . perm ) log_scale = ( self . log_scale @ y ) . permute ( * self . perm ) if self . temperature is not None : log_scale = np . log ( self . temperature ) + log_scale z = loc + torch . exp ( log_scale ) * eps log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow ( eps , 2 ), list ( range ( 1 , self . n_dim + 1 )) ) return z , log_p def log_prob ( self , z , y ): if y . dim () == 1 : y_onehot = torch . zeros ( ( self . num_classes , len ( y )), dtype = self . loc . dtype , device = self . loc . device ) y_onehot . scatter_ ( 0 , y [ None ], 1 ) y = y_onehot else : y = y . t () loc = ( self . loc @ y ) . permute ( * self . perm ) log_scale = ( self . log_scale @ y ) . permute ( * self . perm ) if self . temperature is not None : log_scale = np . log ( self . temperature ) + log_scale log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow (( z - loc ) / torch . exp ( log_scale ), 2 ), list ( range ( 1 , self . n_dim + 1 )), ) return log_p","title":"ClassCondDiagGaussian"},{"location":"references/#normflows.distributions.base.ClassCondDiagGaussian.__init__","text":"Constructor Parameters: Name Type Description Default shape Tuple with shape of data, if int shape has one dimension required num_classes Number of classes to condition on required Source code in normflows/distributions/base.py 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 def __init__ ( self , shape , num_classes ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension num_classes: Number of classes to condition on \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . perm = [ self . n_dim ] + list ( range ( self . n_dim )) self . d = np . prod ( shape ) self . num_classes = num_classes self . loc = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . log_scale = nn . Parameter ( torch . zeros ( * self . shape , num_classes )) self . temperature = None # Temperature parameter for annealed sampling","title":"__init__()"},{"location":"references/#normflows.distributions.base.DiagGaussian","text":"Bases: BaseDistribution Multivariate Gaussian distribution with diagonal covariance matrix Source code in normflows/distributions/base.py 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 class DiagGaussian ( BaseDistribution ): \"\"\" Multivariate Gaussian distribution with diagonal covariance matrix \"\"\" def __init__ ( self , shape , trainable = True ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) if trainable : self . loc = nn . Parameter ( torch . zeros ( 1 , * self . shape )) self . log_scale = nn . Parameter ( torch . zeros ( 1 , * self . shape )) else : self . register_buffer ( \"loc\" , torch . zeros ( 1 , * self . shape )) self . register_buffer ( \"log_scale\" , torch . zeros ( 1 , * self . shape )) self . temperature = None # Temperature parameter for annealed sampling def forward ( self , num_samples = 1 ): eps = torch . randn ( ( num_samples ,) + self . shape , dtype = self . loc . dtype , device = self . loc . device ) if self . temperature is None : log_scale = self . log_scale else : log_scale = self . log_scale + np . log ( self . temperature ) z = self . loc + torch . exp ( log_scale ) * eps log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow ( eps , 2 ), list ( range ( 1 , self . n_dim + 1 )) ) return z , log_p def log_prob ( self , z ): if self . temperature is None : log_scale = self . log_scale else : log_scale = self . log_scale + np . log ( self . temperature ) log_p = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( log_scale + 0.5 * torch . pow (( z - self . loc ) / torch . exp ( log_scale ), 2 ), list ( range ( 1 , self . n_dim + 1 )), ) return log_p","title":"DiagGaussian"},{"location":"references/#normflows.distributions.base.DiagGaussian.__init__","text":"Constructor Parameters: Name Type Description Default shape Tuple with shape of data, if int shape has one dimension required Source code in normflows/distributions/base.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 def __init__ ( self , shape , trainable = True ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . d = np . prod ( shape ) if trainable : self . loc = nn . Parameter ( torch . zeros ( 1 , * self . shape )) self . log_scale = nn . Parameter ( torch . zeros ( 1 , * self . shape )) else : self . register_buffer ( \"loc\" , torch . zeros ( 1 , * self . shape )) self . register_buffer ( \"log_scale\" , torch . zeros ( 1 , * self . shape )) self . temperature = None # Temperature parameter for annealed sampling","title":"__init__()"},{"location":"references/#normflows.distributions.base.GaussianMixture","text":"Bases: BaseDistribution Mixture of Gaussians with diagonal covariance matrix Source code in normflows/distributions/base.py 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 class GaussianMixture ( BaseDistribution ): \"\"\" Mixture of Gaussians with diagonal covariance matrix \"\"\" def __init__ ( self , n_modes , dim , loc = None , scale = None , weights = None , trainable = True ): \"\"\"Constructor Args: n_modes: Number of modes of the mixture model dim: Number of dimensions of each Gaussian loc: List of mean values scale: List of diagonals of the covariance matrices weights: List of mode probabilities trainable: Flag, if true parameters will be optimized during training \"\"\" super () . __init__ () self . n_modes = n_modes self . dim = dim if loc is None : loc = np . random . randn ( self . n_modes , self . dim ) loc = np . array ( loc )[ None , ... ] if scale is None : scale = np . ones (( self . n_modes , self . dim )) scale = np . array ( scale )[ None , ... ] if weights is None : weights = np . ones ( self . n_modes ) weights = np . array ( weights )[ None , ... ] weights /= weights . sum ( 1 ) if trainable : self . loc = nn . Parameter ( torch . tensor ( 1.0 * loc )) self . log_scale = nn . Parameter ( torch . tensor ( np . log ( 1.0 * scale ))) self . weight_scores = nn . Parameter ( torch . tensor ( np . log ( 1.0 * weights ))) else : self . register_buffer ( \"loc\" , torch . tensor ( 1.0 * loc )) self . register_buffer ( \"log_scale\" , torch . tensor ( np . log ( 1.0 * scale ))) self . register_buffer ( \"weight_scores\" , torch . tensor ( np . log ( 1.0 * weights ))) def forward ( self , num_samples = 1 ): # Get weights weights = torch . softmax ( self . weight_scores , 1 ) # Sample mode indices mode = torch . multinomial ( weights [ 0 , :], num_samples , replacement = True ) mode_1h = nn . functional . one_hot ( mode , self . n_modes ) mode_1h = mode_1h [ ... , None ] # Get samples eps_ = torch . randn ( num_samples , self . dim , dtype = self . loc . dtype , device = self . loc . device ) scale_sample = torch . sum ( torch . exp ( self . log_scale ) * mode_1h , 1 ) loc_sample = torch . sum ( self . loc * mode_1h , 1 ) z = eps_ * scale_sample + loc_sample # Compute log probability eps = ( z [:, None , :] - self . loc ) / torch . exp ( self . log_scale ) log_p = ( - 0.5 * self . dim * np . log ( 2 * np . pi ) + torch . log ( weights ) - 0.5 * torch . sum ( torch . pow ( eps , 2 ), 2 ) - torch . sum ( self . log_scale , 2 ) ) log_p = torch . logsumexp ( log_p , 1 ) return z , log_p def log_prob ( self , z ): # Get weights weights = torch . softmax ( self . weight_scores , 1 ) # Compute log probability eps = ( z [:, None , :] - self . loc ) / torch . exp ( self . log_scale ) log_p = ( - 0.5 * self . dim * np . log ( 2 * np . pi ) + torch . log ( weights ) - 0.5 * torch . sum ( torch . pow ( eps , 2 ), 2 ) - torch . sum ( self . log_scale , 2 ) ) log_p = torch . logsumexp ( log_p , 1 ) return log_p","title":"GaussianMixture"},{"location":"references/#normflows.distributions.base.GaussianMixture.__init__","text":"Constructor Parameters: Name Type Description Default n_modes Number of modes of the mixture model required dim Number of dimensions of each Gaussian required loc List of mean values None scale List of diagonals of the covariance matrices None weights List of mode probabilities None trainable Flag, if true parameters will be optimized during training True Source code in normflows/distributions/base.py 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 def __init__ ( self , n_modes , dim , loc = None , scale = None , weights = None , trainable = True ): \"\"\"Constructor Args: n_modes: Number of modes of the mixture model dim: Number of dimensions of each Gaussian loc: List of mean values scale: List of diagonals of the covariance matrices weights: List of mode probabilities trainable: Flag, if true parameters will be optimized during training \"\"\" super () . __init__ () self . n_modes = n_modes self . dim = dim if loc is None : loc = np . random . randn ( self . n_modes , self . dim ) loc = np . array ( loc )[ None , ... ] if scale is None : scale = np . ones (( self . n_modes , self . dim )) scale = np . array ( scale )[ None , ... ] if weights is None : weights = np . ones ( self . n_modes ) weights = np . array ( weights )[ None , ... ] weights /= weights . sum ( 1 ) if trainable : self . loc = nn . Parameter ( torch . tensor ( 1.0 * loc )) self . log_scale = nn . Parameter ( torch . tensor ( np . log ( 1.0 * scale ))) self . weight_scores = nn . Parameter ( torch . tensor ( np . log ( 1.0 * weights ))) else : self . register_buffer ( \"loc\" , torch . tensor ( 1.0 * loc )) self . register_buffer ( \"log_scale\" , torch . tensor ( np . log ( 1.0 * scale ))) self . register_buffer ( \"weight_scores\" , torch . tensor ( np . log ( 1.0 * weights )))","title":"__init__()"},{"location":"references/#normflows.distributions.base.GaussianPCA","text":"Bases: BaseDistribution Gaussian distribution resulting from linearly mapping a normal distributed latent variable describing the \"content of the target\" Source code in normflows/distributions/base.py 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 class GaussianPCA ( BaseDistribution ): \"\"\" Gaussian distribution resulting from linearly mapping a normal distributed latent variable describing the \"content of the target\" \"\"\" def __init__ ( self , dim , latent_dim = None , sigma = 0.1 ): \"\"\"Constructor Args: dim: Number of dimensions of the flow variables latent_dim: Number of dimensions of the latent \"content\" variable; if None it is set equal to dim sigma: Noise level \"\"\" super () . __init__ () self . dim = dim if latent_dim is None : self . latent_dim = dim else : self . latent_dim = latent_dim self . loc = nn . Parameter ( torch . zeros ( 1 , dim )) self . W = nn . Parameter ( torch . randn ( latent_dim , dim )) self . log_sigma = nn . Parameter ( torch . tensor ( np . log ( sigma ))) def forward ( self , num_samples = 1 ): eps = torch . randn ( num_samples , self . latent_dim , dtype = self . loc . dtype , device = self . loc . device ) z_ = torch . matmul ( eps , self . W ) z = z_ + self . loc Sig = torch . matmul ( self . W . T , self . W ) + torch . exp ( self . log_sigma * 2 ) * torch . eye ( self . dim , dtype = self . loc . dtype , device = self . loc . device ) log_p = ( self . dim / 2 * np . log ( 2 * np . pi ) - 0.5 * torch . det ( Sig ) - 0.5 * torch . sum ( z_ * torch . matmul ( z_ , torch . inverse ( Sig )), 1 ) ) return z , log_p def log_prob ( self , z ): z_ = z - self . loc Sig = torch . matmul ( self . W . T , self . W ) + torch . exp ( self . log_sigma * 2 ) * torch . eye ( self . dim , dtype = self . loc . dtype , device = self . loc . device ) log_p = ( self . dim / 2 * np . log ( 2 * np . pi ) - 0.5 * torch . det ( Sig ) - 0.5 * torch . sum ( z_ * torch . matmul ( z_ , torch . inverse ( Sig )), 1 ) ) return log_p","title":"GaussianPCA"},{"location":"references/#normflows.distributions.base.GaussianPCA.__init__","text":"Constructor Parameters: Name Type Description Default dim Number of dimensions of the flow variables required latent_dim Number of dimensions of the latent \"content\" variable; if None it is set equal to dim None sigma Noise level 0.1 Source code in normflows/distributions/base.py 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 def __init__ ( self , dim , latent_dim = None , sigma = 0.1 ): \"\"\"Constructor Args: dim: Number of dimensions of the flow variables latent_dim: Number of dimensions of the latent \"content\" variable; if None it is set equal to dim sigma: Noise level \"\"\" super () . __init__ () self . dim = dim if latent_dim is None : self . latent_dim = dim else : self . latent_dim = latent_dim self . loc = nn . Parameter ( torch . zeros ( 1 , dim )) self . W = nn . Parameter ( torch . randn ( latent_dim , dim )) self . log_sigma = nn . Parameter ( torch . tensor ( np . log ( sigma )))","title":"__init__()"},{"location":"references/#normflows.distributions.base.GlowBase","text":"Bases: BaseDistribution Base distribution of the Glow model, i.e. Diagonal Gaussian with one mean and log scale for each channel Source code in normflows/distributions/base.py 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 class GlowBase ( BaseDistribution ): \"\"\" Base distribution of the Glow model, i.e. Diagonal Gaussian with one mean and log scale for each channel \"\"\" def __init__ ( self , shape , num_classes = None , logscale_factor = 3.0 ): \"\"\"Constructor Args: shape: Shape of the variables num_classes: Number of classes if the base is class conditional, None otherwise logscale_factor: Scaling factor for mean and log variance \"\"\" super () . __init__ () # Save shape and related statistics if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . num_pix = np . prod ( shape [ 1 :]) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . num_classes = num_classes self . class_cond = num_classes is not None self . logscale_factor = logscale_factor # Set up parameters self . loc = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . loc_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) # Class conditional parameter if needed if self . class_cond : self . loc_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ])) self . log_scale_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ]) ) # Temperature parameter for annealed sampling self . temperature = None def forward ( self , num_samples = 1 , y = None ): # Prepare parameter loc = self . loc * torch . exp ( self . loc_logs * self . logscale_factor ) log_scale = self . log_scale * torch . exp ( self . log_scale_logs * self . logscale_factor ) if self . class_cond : if y is not None : num_samples = len ( y ) else : y = torch . randint ( self . num_classes , ( num_samples ,), device = self . loc . device ) if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = self . loc . dtype , device = self . loc . device , ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot loc = loc + ( y @ self . loc_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) log_scale = log_scale + ( y @ self . log_scale_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) if self . temperature is not None : log_scale = log_scale + np . log ( self . temperature ) # Sample eps = torch . randn ( ( num_samples ,) + self . shape , dtype = self . loc . dtype , device = self . loc . device ) z = loc + torch . exp ( log_scale ) * eps # Get log prob log_p = ( - 0.5 * self . d * np . log ( 2 * np . pi ) - self . num_pix * torch . sum ( log_scale , dim = self . sum_dim ) - 0.5 * torch . sum ( torch . pow ( eps , 2 ), dim = self . sum_dim ) ) return z , log_p def log_prob ( self , z , y = None ): # Perpare parameter loc = self . loc * torch . exp ( self . loc_logs * self . logscale_factor ) log_scale = self . log_scale * torch . exp ( self . log_scale_logs * self . logscale_factor ) if self . class_cond : if y . dim () == 1 : y_onehot = torch . zeros ( ( len ( y ), self . num_classes ), dtype = self . loc . dtype , device = self . loc . device , ) y_onehot . scatter_ ( 1 , y [:, None ], 1 ) y = y_onehot loc = loc + ( y @ self . loc_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) log_scale = log_scale + ( y @ self . log_scale_cc ) . view ( y . size ( 0 ), self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ]) ) if self . temperature is not None : log_scale = log_scale + np . log ( self . temperature ) # Get log prob log_p = ( - 0.5 * self . d * np . log ( 2 * np . pi ) - self . num_pix * torch . sum ( log_scale , dim = self . sum_dim ) - 0.5 * torch . sum ( torch . pow (( z - loc ) / torch . exp ( log_scale ), 2 ), dim = self . sum_dim ) ) return log_p","title":"GlowBase"},{"location":"references/#normflows.distributions.base.GlowBase.__init__","text":"Constructor Parameters: Name Type Description Default shape Shape of the variables required num_classes Number of classes if the base is class conditional, None otherwise None logscale_factor Scaling factor for mean and log variance 3.0 Source code in normflows/distributions/base.py 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 def __init__ ( self , shape , num_classes = None , logscale_factor = 3.0 ): \"\"\"Constructor Args: shape: Shape of the variables num_classes: Number of classes if the base is class conditional, None otherwise logscale_factor: Scaling factor for mean and log variance \"\"\" super () . __init__ () # Save shape and related statistics if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . n_dim = len ( shape ) self . num_pix = np . prod ( shape [ 1 :]) self . d = np . prod ( shape ) self . sum_dim = list ( range ( 1 , self . n_dim + 1 )) self . num_classes = num_classes self . class_cond = num_classes is not None self . logscale_factor = logscale_factor # Set up parameters self . loc = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . loc_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) self . log_scale_logs = nn . Parameter ( torch . zeros ( 1 , self . shape [ 0 ], * (( self . n_dim - 1 ) * [ 1 ])) ) # Class conditional parameter if needed if self . class_cond : self . loc_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ])) self . log_scale_cc = nn . Parameter ( torch . zeros ( self . num_classes , self . shape [ 0 ]) ) # Temperature parameter for annealed sampling self . temperature = None","title":"__init__()"},{"location":"references/#normflows.distributions.base.Uniform","text":"Bases: BaseDistribution Multivariate uniform distribution Source code in normflows/distributions/base.py 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 class Uniform ( BaseDistribution ): \"\"\" Multivariate uniform distribution \"\"\" def __init__ ( self , shape , low =- 1.0 , high = 1.0 ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension low: Lower bound of uniform distribution high: Upper bound of uniform distribution \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . d = np . prod ( shape ) self . low = torch . tensor ( low ) self . high = torch . tensor ( high ) self . log_prob_val = - self . d * np . log ( self . high - self . low ) def forward ( self , num_samples = 1 ): eps = torch . rand ( ( num_samples ,) + self . shape , dtype = self . low . dtype , device = self . low . device ) z = self . low + ( self . high - self . low ) * eps log_p = self . log_prob_val * torch . ones ( num_samples , device = self . low . device ) return z , log_p def log_prob ( self , z ): log_p = self . log_prob_val * torch . ones ( z . shape [ 0 ], device = z . device ) out_range = torch . logical_or ( z < self . low , z > self . high ) ind_inf = torch . any ( torch . reshape ( out_range , ( z . shape [ 0 ], - 1 )), dim =- 1 ) log_p [ ind_inf ] = - np . inf return log_p","title":"Uniform"},{"location":"references/#normflows.distributions.base.Uniform.__init__","text":"Constructor Parameters: Name Type Description Default shape Tuple with shape of data, if int shape has one dimension required low Lower bound of uniform distribution -1.0 high Upper bound of uniform distribution 1.0 Source code in normflows/distributions/base.py 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 def __init__ ( self , shape , low =- 1.0 , high = 1.0 ): \"\"\"Constructor Args: shape: Tuple with shape of data, if int shape has one dimension low: Lower bound of uniform distribution high: Upper bound of uniform distribution \"\"\" super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) if isinstance ( shape , list ): shape = tuple ( shape ) self . shape = shape self . d = np . prod ( shape ) self . low = torch . tensor ( low ) self . high = torch . tensor ( high ) self . log_prob_val = - self . d * np . log ( self . high - self . low )","title":"__init__()"},{"location":"references/#normflows.distributions.base.UniformGaussian","text":"Bases: BaseDistribution Distribution of a 1D random variable with some entries having a uniform and others a Gaussian distribution Source code in normflows/distributions/base.py 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 class UniformGaussian ( BaseDistribution ): \"\"\" Distribution of a 1D random variable with some entries having a uniform and others a Gaussian distribution \"\"\" def __init__ ( self , ndim , ind , scale = None ): \"\"\"Constructor Args: ndim: Int, number of dimensions ind: Iterable, indices of uniformly distributed entries scale: Iterable, standard deviation of Gaussian or width of uniform distribution \"\"\" super () . __init__ () self . ndim = ndim if isinstance ( ind , int ): ind = [ ind ] # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) if scale is None : self . register_buffer ( \"scale\" , torch . ones ( self . ndim )) else : self . register_buffer ( \"scale\" , scale ) def forward ( self , num_samples = 1 ): z = self . sample ( num_samples ) return z , self . log_prob ( z ) def sample ( self , num_samples = 1 ): eps_u = ( torch . rand ( ( num_samples , len ( self . ind )), dtype = self . scale . dtype , device = self . scale . device , ) - 0.5 ) eps_g = torch . randn ( ( num_samples , len ( self . ind_ )), dtype = self . scale . dtype , device = self . scale . device , ) z = torch . cat (( eps_u , eps_g ), - 1 ) z = z [ ... , self . inv_perm ] return self . scale * z def log_prob ( self , z ): log_p_u = torch . broadcast_to ( - torch . log ( self . scale [ self . ind ]), ( len ( z ), - 1 )) log_p_g = ( - 0.5 * np . log ( 2 * np . pi ) - torch . log ( self . scale [ self . ind_ ]) - 0.5 * torch . pow ( z [ ... , self . ind_ ] / self . scale [ self . ind_ ], 2 ) ) return torch . sum ( log_p_u , - 1 ) + torch . sum ( log_p_g , - 1 )","title":"UniformGaussian"},{"location":"references/#normflows.distributions.base.UniformGaussian.__init__","text":"Constructor Parameters: Name Type Description Default ndim Int, number of dimensions required ind Iterable, indices of uniformly distributed entries required scale Iterable, standard deviation of Gaussian or width of uniform distribution None Source code in normflows/distributions/base.py 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 def __init__ ( self , ndim , ind , scale = None ): \"\"\"Constructor Args: ndim: Int, number of dimensions ind: Iterable, indices of uniformly distributed entries scale: Iterable, standard deviation of Gaussian or width of uniform distribution \"\"\" super () . __init__ () self . ndim = ndim if isinstance ( ind , int ): ind = [ ind ] # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) if scale is None : self . register_buffer ( \"scale\" , torch . ones ( self . ndim )) else : self . register_buffer ( \"scale\" , scale )","title":"__init__()"},{"location":"references/#normflows.distributions.base_test","text":"","title":"base_test"},{"location":"references/#normflows.distributions.decoder","text":"","title":"decoder"},{"location":"references/#normflows.distributions.decoder.BaseDecoder","text":"Bases: nn . Module Source code in normflows/distributions/decoder.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 class BaseDecoder ( nn . Module ): def __init__ ( self ): super () . __init__ () def forward ( self , z ): \"\"\"Decodes z to x Args: z: latent variable Returns: x, std of x \"\"\" raise NotImplementedError def log_prob ( self , x , z ): \"\"\"Log probability Args: x: observable z: latent variable Returns: log(p) of x given z \"\"\" raise NotImplementedError","title":"BaseDecoder"},{"location":"references/#normflows.distributions.decoder.BaseDecoder.forward","text":"Decodes z to x Parameters: Name Type Description Default z latent variable required Returns: Type Description x, std of x Source code in normflows/distributions/decoder.py 10 11 12 13 14 15 16 17 18 19 def forward ( self , z ): \"\"\"Decodes z to x Args: z: latent variable Returns: x, std of x \"\"\" raise NotImplementedError","title":"forward()"},{"location":"references/#normflows.distributions.decoder.BaseDecoder.log_prob","text":"Log probability Parameters: Name Type Description Default x observable required z latent variable required Returns: Type Description log(p) of x given z Source code in normflows/distributions/decoder.py 21 22 23 24 25 26 27 28 29 30 31 def log_prob ( self , x , z ): \"\"\"Log probability Args: x: observable z: latent variable Returns: log(p) of x given z \"\"\" raise NotImplementedError","title":"log_prob()"},{"location":"references/#normflows.distributions.decoder.NNBernoulliDecoder","text":"Bases: BaseDecoder BaseDecoder representing a Bernoulli distribution with mean parametrized by a NN Source code in normflows/distributions/decoder.py 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 class NNBernoulliDecoder ( BaseDecoder ): \"\"\" BaseDecoder representing a Bernoulli distribution with mean parametrized by a NN \"\"\" def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) \"\"\" super () . __init__ () self . net = net def forward ( self , z ): mean = torch . sigmoid ( self . net ( z )) return mean def log_prob ( self , x , z ): score = self . net ( z ) if len ( z ) > len ( x ): x = x . unsqueeze ( 1 ) x = x . repeat ( 1 , z . size ()[ 0 ] // x . size ()[ 0 ], * (( x . dim () - 2 ) * [ 1 ])) . view ( - 1 , * x . size ()[ 2 :] ) log_sig = lambda a : - torch . relu ( - a ) - torch . log ( 1 + torch . exp ( - torch . abs ( a ))) log_p = torch . sum ( x * log_sig ( score ) + ( 1 - x ) * log_sig ( - score ), list ( range ( 1 , x . dim ())) ) return log_p","title":"NNBernoulliDecoder"},{"location":"references/#normflows.distributions.decoder.NNBernoulliDecoder.__init__","text":"Constructor Parameters: Name Type Description Default net neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) required Source code in normflows/distributions/decoder.py 78 79 80 81 82 83 84 85 def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean Bernoulli (mean = sigmoid(nn_out) \"\"\" super () . __init__ () self . net = net","title":"__init__()"},{"location":"references/#normflows.distributions.decoder.NNDiagGaussianDecoder","text":"Bases: BaseDecoder BaseDecoder representing a diagonal Gaussian distribution with mean and std parametrized by a NN Source code in normflows/distributions/decoder.py 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 class NNDiagGaussianDecoder ( BaseDecoder ): \"\"\" BaseDecoder representing a diagonal Gaussian distribution with mean and std parametrized by a NN \"\"\" def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean and standard deviation of diagonal Gaussian \"\"\" super () . __init__ () self . net = net def forward ( self , z ): mean_std = self . net ( z ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] std = torch . exp ( 0.5 * mean_std [:, n_hidden :, ... ]) return mean , std def log_prob ( self , x , z ): mean_std = self . net ( z ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] var = torch . exp ( mean_std [:, n_hidden :, ... ]) if len ( z ) > len ( x ): x = x . unsqueeze ( 1 ) x = x . repeat ( 1 , z . size ()[ 0 ] // x . size ()[ 0 ], * (( x . dim () - 2 ) * [ 1 ])) . view ( - 1 , * x . size ()[ 2 :] ) log_p = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 1 :])) * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . log ( var ) + ( x - mean ) ** 2 / var , list ( range ( 1 , z . dim ())) ) return log_p","title":"NNDiagGaussianDecoder"},{"location":"references/#normflows.distributions.decoder.NNDiagGaussianDecoder.__init__","text":"Constructor Parameters: Name Type Description Default net neural network parametrizing mean and standard deviation of diagonal Gaussian required Source code in normflows/distributions/decoder.py 39 40 41 42 43 44 45 46 def __init__ ( self , net ): \"\"\"Constructor Args: net: neural network parametrizing mean and standard deviation of diagonal Gaussian \"\"\" super () . __init__ () self . net = net","title":"__init__()"},{"location":"references/#normflows.distributions.decoder_test","text":"","title":"decoder_test"},{"location":"references/#normflows.distributions.distribution_test","text":"","title":"distribution_test"},{"location":"references/#normflows.distributions.distribution_test.DistributionTest","text":"Bases: unittest . TestCase Generic test case for distribution modules Source code in normflows/distributions/distribution_test.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 class DistributionTest ( unittest . TestCase ): \"\"\" Generic test case for distribution modules \"\"\" def assertClose ( self , actual , expected , atol = None , rtol = None ): assert_close ( actual , expected , atol = atol , rtol = rtol ) def checkForward ( self , distribution , num_samples = 1 , ** kwargs ): # Do forward outputs , log_p = distribution ( num_samples , ** kwargs ) # Check type assert outputs . dtype == log_p . dtype # Check shape assert log_p . shape [ 0 ] == num_samples assert outputs . shape [ 0 ] == num_samples # Check dim assert outputs . dim () > log_p . dim () # Return results return outputs , log_p def checkLogProb ( self , distribution , inputs , ** kwargs ): # Compute log prob log_p = distribution . log_prob ( inputs , ** kwargs ) # Check type assert log_p . dtype == inputs . dtype # Check shape assert log_p . shape [ 0 ] == inputs . shape [ 0 ] # Return results return log_p def checkSample ( self , distribution , num_samples = 1 , ** kwargs ): # Do forward outputs = distribution . sample ( num_samples , ** kwargs ) # Check shape assert outputs . shape [ 0 ] == num_samples # Check dim assert outputs . dim () > 1 # Return results return outputs def checkForwardLogProb ( self , distribution , num_samples = 1 , atol = None , rtol = None , ** kwargs ): # Check forward outputs , log_p = self . checkForward ( distribution , num_samples , ** kwargs ) # Check log prob log_p_ = self . checkLogProb ( distribution , outputs , ** kwargs ) # Check consistency self . assertClose ( log_p_ , log_p , atol , rtol )","title":"DistributionTest"},{"location":"references/#normflows.distributions.encoder","text":"","title":"encoder"},{"location":"references/#normflows.distributions.encoder.BaseEncoder","text":"Bases: nn . Module Base distribution of a flow-based variational autoencoder Parameters of the distribution depend of the target variable x Source code in normflows/distributions/encoder.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 class BaseEncoder ( nn . Module ): \"\"\" Base distribution of a flow-based variational autoencoder Parameters of the distribution depend of the target variable x \"\"\" def __init__ ( self ): super () . __init__ () def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, first dimension is batch size num_samples: number of samples to draw per element of mini-batch Returns sample of z for x, log probability for sample \"\"\" raise NotImplementedError def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch size x: Variable to condition on, first dimension is batch size Returns: log probability of z given x \"\"\" raise NotImplementedError","title":"BaseEncoder"},{"location":"references/#normflows.distributions.encoder.BaseEncoder.forward","text":"Parameters: Name Type Description Default x Variable to condition on, first dimension is batch size required num_samples number of samples to draw per element of mini-batch 1 Returns sample of z for x, log probability for sample Source code in normflows/distributions/encoder.py 15 16 17 18 19 20 21 22 23 24 def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, first dimension is batch size num_samples: number of samples to draw per element of mini-batch Returns sample of z for x, log probability for sample \"\"\" raise NotImplementedError","title":"forward()"},{"location":"references/#normflows.distributions.encoder.BaseEncoder.log_prob","text":"Parameters: Name Type Description Default z Primary random variable, first dimension is batch size required x Variable to condition on, first dimension is batch size required Returns: Type Description log probability of z given x Source code in normflows/distributions/encoder.py 26 27 28 29 30 31 32 33 34 35 36 def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch size x: Variable to condition on, first dimension is batch size Returns: log probability of z given x \"\"\" raise NotImplementedError","title":"log_prob()"},{"location":"references/#normflows.distributions.encoder.ConstDiagGaussian","text":"Bases: BaseEncoder Source code in normflows/distributions/encoder.py 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 class ConstDiagGaussian ( BaseEncoder ): def __init__ ( self , loc , scale ): \"\"\"Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x Args: loc: mean vector of the distribution scale: vector of the standard deviations on the diagonal of the covariance matrix \"\"\" super () . __init__ () self . d = len ( loc ) if not torch . is_tensor ( loc ): loc = torch . tensor ( loc ) if not torch . is_tensor ( scale ): scale = torch . tensor ( scale ) self . loc = nn . Parameter ( loc . reshape (( 1 , 1 , self . d ))) self . scale = nn . Parameter ( scale ) def forward ( self , x = None , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, will only be used to determine the batch size num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" if x is not None : batch_size = len ( x ) else : batch_size = 1 eps = torch . randn (( batch_size , num_samples , self . d ), device = x . device ) z = self . loc + self . scale * eps log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * torch . pow ( eps , 2 ), 2 ) return z , log_q def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * (( z - self . loc ) / self . scale ) ** 2 , 2 ) return log_q","title":"ConstDiagGaussian"},{"location":"references/#normflows.distributions.encoder.ConstDiagGaussian.__init__","text":"Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x Parameters: Name Type Description Default loc mean vector of the distribution required scale vector of the standard deviations on the diagonal of the covariance matrix required Source code in normflows/distributions/encoder.py 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 def __init__ ( self , loc , scale ): \"\"\"Multivariate Gaussian distribution with diagonal covariance and parameters being constant wrt x Args: loc: mean vector of the distribution scale: vector of the standard deviations on the diagonal of the covariance matrix \"\"\" super () . __init__ () self . d = len ( loc ) if not torch . is_tensor ( loc ): loc = torch . tensor ( loc ) if not torch . is_tensor ( scale ): scale = torch . tensor ( scale ) self . loc = nn . Parameter ( loc . reshape (( 1 , 1 , self . d ))) self . scale = nn . Parameter ( scale )","title":"__init__()"},{"location":"references/#normflows.distributions.encoder.ConstDiagGaussian.forward","text":"Parameters: Name Type Description Default x Variable to condition on, will only be used to determine the batch size None num_samples number of samples to draw per element of mini-batch 1 Returns: Type Description sample of z for x, log probability for sample Source code in normflows/distributions/encoder.py 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 def forward ( self , x = None , num_samples = 1 ): \"\"\" Args: x: Variable to condition on, will only be used to determine the batch size num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" if x is not None : batch_size = len ( x ) else : batch_size = 1 eps = torch . randn (( batch_size , num_samples , self . d ), device = x . device ) z = self . loc + self . scale * eps log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * torch . pow ( eps , 2 ), 2 ) return z , log_q","title":"forward()"},{"location":"references/#normflows.distributions.encoder.ConstDiagGaussian.log_prob","text":"Parameters: Name Type Description Default z Primary random variable, first dimension is batch dimension required x Variable to condition on, first dimension is batch dimension required Returns: Type Description log probability of z given x Source code in normflows/distributions/encoder.py 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) log_q = - 0.5 * self . d * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * (( z - self . loc ) / self . scale ) ** 2 , 2 ) return log_q","title":"log_prob()"},{"location":"references/#normflows.distributions.encoder.NNDiagGaussian","text":"Bases: BaseEncoder Diagonal Gaussian distribution with mean and variance determined by a neural network Source code in normflows/distributions/encoder.py 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 class NNDiagGaussian ( BaseEncoder ): \"\"\" Diagonal Gaussian distribution with mean and variance determined by a neural network \"\"\" def __init__ ( self , net ): \"\"\"Construtor Args: net: net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) \"\"\" super () . __init__ () self . net = net def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" batch_size = len ( x ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) std = torch . exp ( 0.5 * mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) eps = torch . randn ( ( batch_size , num_samples ) + tuple ( mean . size ()[ 2 :]), device = x . device ) z = mean + std * eps log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( std ) + 0.5 * torch . pow ( eps , 2 ), list ( range ( 2 , z . dim ()))) return z , log_q def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) var = torch . exp ( mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . log ( var ) + ( z - mean ) ** 2 / var , 2 ) return log_q","title":"NNDiagGaussian"},{"location":"references/#normflows.distributions.encoder.NNDiagGaussian.__init__","text":"Construtor Parameters: Name Type Description Default net net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) required Source code in normflows/distributions/encoder.py 135 136 137 138 139 140 141 142 def __init__ ( self , net ): \"\"\"Construtor Args: net: net computing mean (first n / 2 outputs), standard deviation (second n / 2 outputs) \"\"\" super () . __init__ () self . net = net","title":"__init__()"},{"location":"references/#normflows.distributions.encoder.NNDiagGaussian.forward","text":"Parameters: Name Type Description Default x Variable to condition on required num_samples number of samples to draw per element of mini-batch 1 Returns: Type Description sample of z for x, log probability for sample Source code in normflows/distributions/encoder.py 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 def forward ( self , x , num_samples = 1 ): \"\"\" Args: x: Variable to condition on num_samples: number of samples to draw per element of mini-batch Returns: sample of z for x, log probability for sample \"\"\" batch_size = len ( x ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) std = torch . exp ( 0.5 * mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) eps = torch . randn ( ( batch_size , num_samples ) + tuple ( mean . size ()[ 2 :]), device = x . device ) z = mean + std * eps log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( std ) + 0.5 * torch . pow ( eps , 2 ), list ( range ( 2 , z . dim ()))) return z , log_q","title":"forward()"},{"location":"references/#normflows.distributions.encoder.NNDiagGaussian.log_prob","text":"Parameters: Name Type Description Default z Primary random variable, first dimension is batch dimension required x Variable to condition on, first dimension is batch dimension required Returns: Type Description log probability of z given x Source code in normflows/distributions/encoder.py 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 def log_prob ( self , z , x ): \"\"\" Args: z: Primary random variable, first dimension is batch dimension x: Variable to condition on, first dimension is batch dimension Returns: log probability of z given x \"\"\" if z . dim () == 1 : z = z . unsqueeze ( 0 ) if z . dim () == 2 : z = z . unsqueeze ( 0 ) mean_std = self . net ( x ) n_hidden = mean_std . size ()[ 1 ] // 2 mean = mean_std [:, : n_hidden , ... ] . unsqueeze ( 1 ) var = torch . exp ( mean_std [:, n_hidden : ( 2 * n_hidden ), ... ] . unsqueeze ( 1 )) log_q = - 0.5 * torch . prod ( torch . tensor ( z . size ()[ 2 :])) * np . log ( 2 * np . pi ) - 0.5 * torch . sum ( torch . log ( var ) + ( z - mean ) ** 2 / var , 2 ) return log_q","title":"log_prob()"},{"location":"references/#normflows.distributions.encoder_test","text":"","title":"encoder_test"},{"location":"references/#normflows.distributions.linear_interpolation","text":"","title":"linear_interpolation"},{"location":"references/#normflows.distributions.linear_interpolation.LinearInterpolation","text":"Linear interpolation of two distributions in the log space Source code in normflows/distributions/linear_interpolation.py 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 class LinearInterpolation : \"\"\" Linear interpolation of two distributions in the log space \"\"\" def __init__ ( self , dist1 , dist2 , alpha ): \"\"\"Constructor Interpolation parameter alpha: ``` log_p = alpha * log_p_1 + (1 - alpha) * log_p_2 ``` Args: dist1: First distribution dist2: Second distribution alpha: Interpolation parameter \"\"\" self . alpha = alpha self . dist1 = dist1 self . dist2 = dist2 def log_prob ( self , z ): return self . alpha * self . dist1 . log_prob ( z ) + ( 1 - self . alpha ) * self . dist2 . log_prob ( z )","title":"LinearInterpolation"},{"location":"references/#normflows.distributions.linear_interpolation.LinearInterpolation.__init__","text":"Constructor Interpolation parameter alpha: log_p = alpha * log_p_1 + (1 - alpha) * log_p_2 Parameters: Name Type Description Default dist1 First distribution required dist2 Second distribution required alpha Interpolation parameter required Source code in normflows/distributions/linear_interpolation.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 def __init__ ( self , dist1 , dist2 , alpha ): \"\"\"Constructor Interpolation parameter alpha: ``` log_p = alpha * log_p_1 + (1 - alpha) * log_p_2 ``` Args: dist1: First distribution dist2: Second distribution alpha: Interpolation parameter \"\"\" self . alpha = alpha self . dist1 = dist1 self . dist2 = dist2","title":"__init__()"},{"location":"references/#normflows.distributions.mh_proposal","text":"","title":"mh_proposal"},{"location":"references/#normflows.distributions.mh_proposal.DiagGaussianProposal","text":"Bases: MHProposal Diagonal Gaussian distribution with previous value as mean as a proposal for Metropolis Hastings algorithm Source code in normflows/distributions/mh_proposal.py 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 class DiagGaussianProposal ( MHProposal ): \"\"\" Diagonal Gaussian distribution with previous value as mean as a proposal for Metropolis Hastings algorithm \"\"\" def __init__ ( self , shape , scale ): \"\"\"Constructor Args: shape: Shape of variables to sample scale: Standard deviation of distribution \"\"\" super () . __init__ () self . shape = shape self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu . unsqueeze ( 0 )) def sample ( self , z ): num_samples = len ( z ) eps = torch . randn (( num_samples ,) + self . shape , dtype = z . dtype , device = z . device ) z_ = eps * self . scale + z return z_ def log_prob ( self , z_ , z ): log_p = - 0.5 * np . prod ( self . shape ) * np . log ( 2 * np . pi ) - torch . sum ( torch . log ( self . scale ) + 0.5 * torch . pow (( z_ - z ) / self . scale , 2 ), list ( range ( 1 , z . dim ())), ) return log_p def forward ( self , z ): num_samples = len ( z ) eps = torch . randn (( num_samples ,) + self . shape , dtype = z . dtype , device = z . device ) z_ = eps * self . scale + z log_p_diff = torch . zeros ( num_samples , dtype = z . dtype , device = z . device ) return z_ , log_p_diff","title":"DiagGaussianProposal"},{"location":"references/#normflows.distributions.mh_proposal.DiagGaussianProposal.__init__","text":"Constructor Parameters: Name Type Description Default shape Shape of variables to sample required scale Standard deviation of distribution required Source code in normflows/distributions/mh_proposal.py 53 54 55 56 57 58 59 60 61 62 63 def __init__ ( self , shape , scale ): \"\"\"Constructor Args: shape: Shape of variables to sample scale: Standard deviation of distribution \"\"\" super () . __init__ () self . shape = shape self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu . unsqueeze ( 0 ))","title":"__init__()"},{"location":"references/#normflows.distributions.mh_proposal.MHProposal","text":"Bases: nn . Module Proposal distribution for the Metropolis Hastings algorithm Source code in normflows/distributions/mh_proposal.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 class MHProposal ( nn . Module ): \"\"\" Proposal distribution for the Metropolis Hastings algorithm \"\"\" def __init__ ( self ): super () . __init__ () def sample ( self , z ): \"\"\" Sample new value based on previous z \"\"\" raise NotImplementedError def log_prob ( self , z_ , z ): \"\"\" Args: z_: Potential new sample z: Previous sample Returns: Log probability of proposal distribution \"\"\" raise NotImplementedError def forward ( self , z ): \"\"\"Draw samples given z and compute log probability difference ``` log(p(z | z_new)) - log(p(z_new | z)) ``` Args: z: Previous samples Returns: Proposal, difference of log probability ratio \"\"\" raise NotImplementedError","title":"MHProposal"},{"location":"references/#normflows.distributions.mh_proposal.MHProposal.forward","text":"Draw samples given z and compute log probability difference log(p(z | z_new)) - log(p(z_new | z)) Parameters: Name Type Description Default z Previous samples required Returns: Type Description Proposal, difference of log probability ratio Source code in normflows/distributions/mh_proposal.py 31 32 33 34 35 36 37 38 39 40 41 42 43 44 def forward ( self , z ): \"\"\"Draw samples given z and compute log probability difference ``` log(p(z | z_new)) - log(p(z_new | z)) ``` Args: z: Previous samples Returns: Proposal, difference of log probability ratio \"\"\" raise NotImplementedError","title":"forward()"},{"location":"references/#normflows.distributions.mh_proposal.MHProposal.log_prob","text":"Parameters: Name Type Description Default z_ Potential new sample required z Previous sample required Returns: Type Description Log probability of proposal distribution Source code in normflows/distributions/mh_proposal.py 20 21 22 23 24 25 26 27 28 29 def log_prob ( self , z_ , z ): \"\"\" Args: z_: Potential new sample z: Previous sample Returns: Log probability of proposal distribution \"\"\" raise NotImplementedError","title":"log_prob()"},{"location":"references/#normflows.distributions.mh_proposal.MHProposal.sample","text":"Sample new value based on previous z Source code in normflows/distributions/mh_proposal.py 14 15 16 17 18 def sample ( self , z ): \"\"\" Sample new value based on previous z \"\"\" raise NotImplementedError","title":"sample()"},{"location":"references/#normflows.distributions.prior","text":"","title":"prior"},{"location":"references/#normflows.distributions.prior.ImagePrior","text":"Bases: nn . Module Intensities of an image determine probability density of prior Source code in normflows/distributions/prior.py 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 class ImagePrior ( nn . Module ): \"\"\" Intensities of an image determine probability density of prior \"\"\" def __init__ ( self , image , x_range = [ - 3 , 3 ], y_range = [ - 3 , 3 ], eps = 1.0e-10 ): \"\"\"Constructor Args: image: image as np matrix x_range: x range to position image at y_range: y range to position image at eps: small value to add to image to avoid log(0) problems \"\"\" super () . __init__ () image_ = np . flip ( image , 0 ) . transpose () + eps self . image_cpu = torch . tensor ( image_ / np . max ( image_ )) self . image_size_cpu = self . image_cpu . size () self . x_range = torch . tensor ( x_range ) self . y_range = torch . tensor ( y_range ) self . register_buffer ( \"image\" , self . image_cpu ) self . register_buffer ( \"image_size\" , torch . tensor ( self . image_size_cpu ) . unsqueeze ( 0 ) ) self . register_buffer ( \"density\" , torch . log ( self . image_cpu / torch . sum ( self . image_cpu )) ) self . register_buffer ( \"scale\" , torch . tensor ( [[ self . x_range [ 1 ] - self . x_range [ 0 ], self . y_range [ 1 ] - self . y_range [ 0 ]]] ), ) self . register_buffer ( \"shift\" , torch . tensor ([[ self . x_range [ 0 ], self . y_range [ 0 ]]]) ) def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" z_ = torch . clamp (( z - self . shift ) / self . scale , max = 1 , min = 0 ) ind = ( z_ * ( self . image_size - 1 )) . long () return self . density [ ind [:, 0 ], ind [:, 1 ]] def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" z_ = torch . rand ( ( num_steps , 2 ), dtype = self . image . dtype , device = self . image . device ) prob = torch . rand ( num_steps , dtype = self . image . dtype , device = self . image . device ) ind = ( z_ * ( self . image_size - 1 )) . long () intensity = self . image [ ind [:, 0 ], ind [:, 1 ]] accept = intensity > prob z = z_ [ accept , :] * self . scale + self . shift return z def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . ones (( 0 , 2 ), dtype = self . image . dtype , device = self . image . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z","title":"ImagePrior"},{"location":"references/#normflows.distributions.prior.ImagePrior.__init__","text":"Constructor Parameters: Name Type Description Default image image as np matrix required x_range x range to position image at [-3, 3] y_range y range to position image at [-3, 3] eps small value to add to image to avoid log(0) problems 1e-10 Source code in normflows/distributions/prior.py 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 def __init__ ( self , image , x_range = [ - 3 , 3 ], y_range = [ - 3 , 3 ], eps = 1.0e-10 ): \"\"\"Constructor Args: image: image as np matrix x_range: x range to position image at y_range: y range to position image at eps: small value to add to image to avoid log(0) problems \"\"\" super () . __init__ () image_ = np . flip ( image , 0 ) . transpose () + eps self . image_cpu = torch . tensor ( image_ / np . max ( image_ )) self . image_size_cpu = self . image_cpu . size () self . x_range = torch . tensor ( x_range ) self . y_range = torch . tensor ( y_range ) self . register_buffer ( \"image\" , self . image_cpu ) self . register_buffer ( \"image_size\" , torch . tensor ( self . image_size_cpu ) . unsqueeze ( 0 ) ) self . register_buffer ( \"density\" , torch . log ( self . image_cpu / torch . sum ( self . image_cpu )) ) self . register_buffer ( \"scale\" , torch . tensor ( [[ self . x_range [ 1 ] - self . x_range [ 0 ], self . y_range [ 1 ] - self . y_range [ 0 ]]] ), ) self . register_buffer ( \"shift\" , torch . tensor ([[ self . x_range [ 0 ], self . y_range [ 0 ]]]) )","title":"__init__()"},{"location":"references/#normflows.distributions.prior.ImagePrior.log_prob","text":"Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 59 60 61 62 63 64 65 66 67 68 69 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" z_ = torch . clamp (( z - self . shift ) / self . scale , max = 1 , min = 0 ) ind = ( z_ * ( self . image_size - 1 )) . long () return self . density [ ind [:, 0 ], ind [:, 1 ]]","title":"log_prob()"},{"location":"references/#normflows.distributions.prior.ImagePrior.rejection_sampling","text":"Perform rejection sampling on image distribution Parameters: Name Type Description Default num_steps Number of rejection sampling steps to perform 1 Returns: Type Description Accepted samples Source code in normflows/distributions/prior.py 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" z_ = torch . rand ( ( num_steps , 2 ), dtype = self . image . dtype , device = self . image . device ) prob = torch . rand ( num_steps , dtype = self . image . dtype , device = self . image . device ) ind = ( z_ * ( self . image_size - 1 )) . long () intensity = self . image [ ind [:, 0 ], ind [:, 1 ]] accept = intensity > prob z = z_ [ accept , :] * self . scale + self . shift return z","title":"rejection_sampling()"},{"location":"references/#normflows.distributions.prior.ImagePrior.sample","text":"Sample from image distribution through rejection sampling Parameters: Name Type Description Default num_samples Number of samples to draw 1 Returns: Type Description Samples Source code in normflows/distributions/prior.py 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . ones (( 0 , 2 ), dtype = self . image . dtype , device = self . image . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z","title":"sample()"},{"location":"references/#normflows.distributions.prior.PriorDistribution","text":"Source code in normflows/distributions/prior.py 6 7 8 9 10 11 12 13 14 15 16 17 18 class PriorDistribution : def __init__ ( self ): raise NotImplementedError def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError","title":"PriorDistribution"},{"location":"references/#normflows.distributions.prior.PriorDistribution.log_prob","text":"Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 10 11 12 13 14 15 16 17 18 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError","title":"log_prob()"},{"location":"references/#normflows.distributions.prior.Sinusoidal","text":"Bases: PriorDistribution Source code in normflows/distributions/prior.py 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 class Sinusoidal ( PriorDistribution ): def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density given by ``` w_1(z) = sin(2*pi / period * z[0]) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 ``` Args: scale: scale of the distribution, see formula period: period of the sinosoidal \"\"\" self . scale = scale self . period = period def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 w_1(z) = sin(2*pi / period * z[0]) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) log_prob = ( - 0.5 * (( z_ [ 1 ] - w_1 ( z_ )) / ( self . scale )) ** 2 - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) # add Gaussian envelope for valid p(z) return log_prob","title":"Sinusoidal"},{"location":"references/#normflows.distributions.prior.Sinusoidal.__init__","text":"Distribution 2d with sinusoidal density given by w_1(z) = sin(2*pi / period * z[0]) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 Parameters: Name Type Description Default scale scale of the distribution, see formula required period period of the sinosoidal required Source code in normflows/distributions/prior.py 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density given by ``` w_1(z) = sin(2*pi / period * z[0]) log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 ``` Args: scale: scale of the distribution, see formula period: period of the sinosoidal \"\"\" self . scale = scale self . period = period","title":"__init__()"},{"location":"references/#normflows.distributions.prior.Sinusoidal.log_prob","text":"log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 w_1(z) = sin(2*pi / period * z[0]) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((z[1] - w_1(z)) / (2 * scale)) ** 2 w_1(z) = sin(2*pi / period * z[0]) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) log_prob = ( - 0.5 * (( z_ [ 1 ] - w_1 ( z_ )) / ( self . scale )) ** 2 - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) # add Gaussian envelope for valid p(z) return log_prob","title":"log_prob()"},{"location":"references/#normflows.distributions.prior.Sinusoidal_gap","text":"Bases: PriorDistribution Source code in normflows/distributions/prior.py 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 class Sinusoidal_gap ( PriorDistribution ): def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with gap given by ``` w_1(z) = sin(2*pi / period * z[0]) w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w2_scale = 0.6 self . w2_amp = 3.0 self . w2_mu = 1.0 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_2 = lambda x : self . w2_amp * torch . exp ( - 0.5 * (( z_ [ 0 ] - self . w2_mu ) / self . w2_scale ) ** 2 ) eps = torch . abs ( w_2 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_2 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / self . scale ) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob","title":"Sinusoidal_gap"},{"location":"references/#normflows.distributions.prior.Sinusoidal_gap.__init__","text":"Distribution 2d with sinusoidal density with gap given by w_1(z) = sin(2*pi / period * z[0]) w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2)) Parameters: Name Type Description Default loc distance of modes from the origin required scale scale of modes required Source code in normflows/distributions/prior.py 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with gap given by ``` w_1(z) = sin(2*pi / period * z[0]) w_2(z) = 3 * exp(-0.5 * ((z[0] - 1) / 0.6) ** 2) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.35) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_2(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w2_scale = 0.6 self . w2_amp = 3.0 self . w2_mu = 1.0","title":"__init__()"},{"location":"references/#normflows.distributions.prior.Sinusoidal_gap.log_prob","text":"Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_2 = lambda x : self . w2_amp * torch . exp ( - 0.5 * (( z_ [ 0 ] - self . w2_mu ) / self . w2_scale ) ** 2 ) eps = torch . abs ( w_2 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_2 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / self . scale ) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob","title":"log_prob()"},{"location":"references/#normflows.distributions.prior.Sinusoidal_split","text":"Bases: PriorDistribution Source code in normflows/distributions/prior.py 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 class Sinusoidal_split ( PriorDistribution ): def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with split given by ``` w_1(z) = sin(2*pi / period * z[0]) w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w3_scale = 0.3 self . w3_amp = 3.0 self . w3_mu = 1.0 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_3 = lambda x : self . w3_amp * torch . sigmoid ( ( z_ [ 0 ] - self . w3_mu ) / self . w3_scale ) eps = torch . abs ( w_3 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_3 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / ( self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob","title":"Sinusoidal_split"},{"location":"references/#normflows.distributions.prior.Sinusoidal_split.__init__","text":"Distribution 2d with sinusoidal density with split given by w_1(z) = sin(2*pi / period * z[0]) w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2)) Parameters: Name Type Description Default loc distance of modes from the origin required scale scale of modes required Source code in normflows/distributions/prior.py 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 def __init__ ( self , scale , period ): \"\"\"Distribution 2d with sinusoidal density with split given by ``` w_1(z) = sin(2*pi / period * z[0]) w_3(z) = 3 * sigmoid((z[0] - 1) / 0.3) log(p) = -log(exp(-0.5 * ((z[1] - w_1(z)) / 0.4) ** 2) + exp(-0.5 * ((z[1] - w_1(z) + w_3(z)) / 0.35) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . scale = scale self . period = period self . w3_scale = 0.3 self . w3_amp = 3.0 self . w3_mu = 1.0","title":"__init__()"},{"location":"references/#normflows.distributions.prior.Sinusoidal_split.log_prob","text":"Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z w_1 = lambda x : torch . sin ( 2 * np . pi / self . period * z_ [ 0 ]) w_3 = lambda x : self . w3_amp * torch . sigmoid ( ( z_ [ 0 ] - self . w3_mu ) / self . w3_scale ) eps = torch . abs ( w_3 ( z_ ) / 2 ) a = torch . abs ( z_ [ 1 ] - w_1 ( z_ ) + w_3 ( z_ ) / 2 ) log_prob = ( - 0.5 * (( a - eps ) / ( self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( eps * a ) / self . scale ** 2 )) - 0.5 * ( torch . norm ( z_ , dim = 0 , p = 4 ) / ( 20 * self . scale )) ** 4 ) return log_prob","title":"log_prob()"},{"location":"references/#normflows.distributions.prior.Smiley","text":"Bases: PriorDistribution Source code in normflows/distributions/prior.py 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 class Smiley ( PriorDistribution ): def __init__ ( self , scale ): \"\"\"Distribution 2d of a smiley :) Args: scale: scale of the smiley \"\"\" self . scale = scale self . loc = 2.0 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z log_prob = ( - 0.5 * (( torch . norm ( z_ , dim = 0 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( torch . abs ( z_ [ 1 ] + 0.8 ) - 1.2 ) / ( 2 * self . scale )) ** 2 ) return log_prob","title":"Smiley"},{"location":"references/#normflows.distributions.prior.Smiley.__init__","text":"Distribution 2d of a smiley :) Parameters: Name Type Description Default scale scale of the smiley required Source code in normflows/distributions/prior.py 300 301 302 303 304 305 306 307 def __init__ ( self , scale ): \"\"\"Distribution 2d of a smiley :) Args: scale: scale of the smiley \"\"\" self . scale = scale self . loc = 2.0","title":"__init__()"},{"location":"references/#normflows.distributions.prior.Smiley.log_prob","text":"Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" if z . dim () > 1 : z_ = z . permute (( z . dim () - 1 ,) + tuple ( range ( 0 , z . dim () - 1 ))) else : z_ = z log_prob = ( - 0.5 * (( torch . norm ( z_ , dim = 0 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( torch . abs ( z_ [ 1 ] + 0.8 ) - 1.2 ) / ( 2 * self . scale )) ** 2 ) return log_prob","title":"log_prob()"},{"location":"references/#normflows.distributions.prior.TwoModes","text":"Bases: PriorDistribution Source code in normflows/distributions/prior.py 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 class TwoModes ( PriorDistribution ): def __init__ ( self , loc , scale ): \"\"\"Distribution 2d with two modes Distribution 2d with two modes at ```z[0] = -loc``` and ```z[0] = loc``` following the density ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . loc = loc self . scale = scale def log_prob ( self , z ): \"\"\" ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) eps = torch . abs ( torch . tensor ( self . loc )) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( a - eps ) / ( 3 * self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( a * eps ) / ( 3 * self . scale ) ** 2 )) ) return log_prob","title":"TwoModes"},{"location":"references/#normflows.distributions.prior.TwoModes.__init__","text":"Distribution 2d with two modes Distribution 2d with two modes at z[0] = -loc and z[0] = loc following the density log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) Args: loc: distance of modes from the origin scale: scale of modes Source code in normflows/distributions/prior.py 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 def __init__ ( self , loc , scale ): \"\"\"Distribution 2d with two modes Distribution 2d with two modes at ```z[0] = -loc``` and ```z[0] = loc``` following the density ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: loc: distance of modes from the origin scale: scale of modes \"\"\" self . loc = loc self . scale = scale","title":"__init__()"},{"location":"references/#normflows.distributions.prior.TwoModes.log_prob","text":"log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/prior.py 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 def log_prob ( self , z ): \"\"\" ``` log(p) = 1/2 * ((norm(z) - loc) / (2 * scale)) ** 2 - log(exp(-1/2 * ((z[0] - loc) / (3 * scale)) ** 2) + exp(-1/2 * ((z[0] + loc) / (3 * scale)) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) eps = torch . abs ( torch . tensor ( self . loc )) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - self . loc ) / ( 2 * self . scale )) ** 2 - 0.5 * (( a - eps ) / ( 3 * self . scale )) ** 2 + torch . log ( 1 + torch . exp ( - 2 * ( a * eps ) / ( 3 * self . scale ) ** 2 )) ) return log_prob","title":"log_prob()"},{"location":"references/#normflows.distributions.prior_test","text":"","title":"prior_test"},{"location":"references/#normflows.distributions.target","text":"","title":"target"},{"location":"references/#normflows.distributions.target.CircularGaussianMixture","text":"Bases: nn . Module Two-dimensional Gaussian mixture arranged in a circle Source code in normflows/distributions/target.py 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 class CircularGaussianMixture ( nn . Module ): \"\"\" Two-dimensional Gaussian mixture arranged in a circle \"\"\" def __init__ ( self , n_modes = 8 ): \"\"\"Constructor Args: n_modes: Number of modes \"\"\" super ( CircularGaussianMixture , self ) . __init__ () self . n_modes = n_modes self . register_buffer ( \"scale\" , torch . tensor ( 2 / 3 * np . sin ( np . pi / self . n_modes )) . float () ) def log_prob ( self , z ): d = torch . zeros (( len ( z ), 0 ), dtype = z . dtype , device = z . device ) for i in range ( self . n_modes ): d_ = ( ( z [:, 0 ] - 2 * np . sin ( 2 * np . pi / self . n_modes * i )) ** 2 + ( z [:, 1 ] - 2 * np . cos ( 2 * np . pi / self . n_modes * i )) ** 2 ) / ( 2 * self . scale ** 2 ) d = torch . cat (( d , d_ [:, None ]), 1 ) log_p = - torch . log ( 2 * np . pi * self . scale ** 2 * self . n_modes ) + torch . logsumexp ( - d , 1 ) return log_p def sample ( self , num_samples = 1 ): eps = torch . randn ( ( num_samples , 2 ), dtype = self . scale . dtype , device = self . scale . device ) phi = ( 2 * np . pi / self . n_modes * torch . randint ( 0 , self . n_modes , ( num_samples ,), device = self . scale . device ) ) loc = torch . stack (( 2 * torch . sin ( phi ), 2 * torch . cos ( phi )), 1 ) . type ( eps . dtype ) return eps * self . scale + loc","title":"CircularGaussianMixture"},{"location":"references/#normflows.distributions.target.CircularGaussianMixture.__init__","text":"Constructor Parameters: Name Type Description Default n_modes Number of modes 8 Source code in normflows/distributions/target.py 137 138 139 140 141 142 143 144 145 146 147 def __init__ ( self , n_modes = 8 ): \"\"\"Constructor Args: n_modes: Number of modes \"\"\" super ( CircularGaussianMixture , self ) . __init__ () self . n_modes = n_modes self . register_buffer ( \"scale\" , torch . tensor ( 2 / 3 * np . sin ( np . pi / self . n_modes )) . float () )","title":"__init__()"},{"location":"references/#normflows.distributions.target.RingMixture","text":"Bases: Target Mixture of ring distributions in two dimensions Source code in normflows/distributions/target.py 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 class RingMixture ( Target ): \"\"\" Mixture of ring distributions in two dimensions \"\"\" def __init__ ( self , n_rings = 2 ): super () . __init__ () self . n_dims = 2 self . max_log_prob = 0.0 self . n_rings = n_rings self . scale = 1 / 4 / self . n_rings def log_prob ( self , z ): d = torch . zeros (( len ( z ), 0 ), dtype = z . dtype , device = z . device ) for i in range ( self . n_rings ): d_ = (( torch . norm ( z , dim = 1 ) - 2 / self . n_rings * ( i + 1 )) ** 2 ) / ( 2 * self . scale ** 2 ) d = torch . cat (( d , d_ [:, None ]), 1 ) return torch . logsumexp ( - d , 1 )","title":"RingMixture"},{"location":"references/#normflows.distributions.target.Target","text":"Bases: nn . Module Sample target distributions to test models Source code in normflows/distributions/target.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 class Target ( nn . Module ): \"\"\" Sample target distributions to test models \"\"\" def __init__ ( self , prop_scale = torch . tensor ( 6.0 ), prop_shift = torch . tensor ( - 3.0 )): \"\"\"Constructor Args: prop_scale: Scale for the uniform proposal prop_shift: Shift for the uniform proposal \"\"\" super () . __init__ () self . register_buffer ( \"prop_scale\" , prop_scale ) self . register_buffer ( \"prop_shift\" , prop_shift ) def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError ( \"The log probability is not implemented yet.\" ) def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" eps = torch . rand ( ( num_steps , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device , ) z_ = self . prop_scale * eps + self . prop_shift prob = torch . rand ( num_steps , dtype = self . prop_scale . dtype , device = self . prop_scale . device ) prob_ = torch . exp ( self . log_prob ( z_ ) - self . max_log_prob ) accept = prob_ > prob z = z_ [ accept , :] return z def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . zeros ( ( 0 , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z","title":"Target"},{"location":"references/#normflows.distributions.target.Target.__init__","text":"Constructor Parameters: Name Type Description Default prop_scale Scale for the uniform proposal torch.tensor(6.0) prop_shift Shift for the uniform proposal torch.tensor(-3.0) Source code in normflows/distributions/target.py 13 14 15 16 17 18 19 20 21 22 def __init__ ( self , prop_scale = torch . tensor ( 6.0 ), prop_shift = torch . tensor ( - 3.0 )): \"\"\"Constructor Args: prop_scale: Scale for the uniform proposal prop_shift: Shift for the uniform proposal \"\"\" super () . __init__ () self . register_buffer ( \"prop_scale\" , prop_scale ) self . register_buffer ( \"prop_shift\" , prop_shift )","title":"__init__()"},{"location":"references/#normflows.distributions.target.Target.log_prob","text":"Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/target.py 24 25 26 27 28 29 30 31 32 def log_prob ( self , z ): \"\"\" Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" raise NotImplementedError ( \"The log probability is not implemented yet.\" )","title":"log_prob()"},{"location":"references/#normflows.distributions.target.Target.rejection_sampling","text":"Perform rejection sampling on image distribution Parameters: Name Type Description Default num_steps Number of rejection sampling steps to perform 1 Returns: Type Description Accepted samples Source code in normflows/distributions/target.py 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def rejection_sampling ( self , num_steps = 1 ): \"\"\"Perform rejection sampling on image distribution Args: num_steps: Number of rejection sampling steps to perform Returns: Accepted samples \"\"\" eps = torch . rand ( ( num_steps , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device , ) z_ = self . prop_scale * eps + self . prop_shift prob = torch . rand ( num_steps , dtype = self . prop_scale . dtype , device = self . prop_scale . device ) prob_ = torch . exp ( self . log_prob ( z_ ) - self . max_log_prob ) accept = prob_ > prob z = z_ [ accept , :] return z","title":"rejection_sampling()"},{"location":"references/#normflows.distributions.target.Target.sample","text":"Sample from image distribution through rejection sampling Parameters: Name Type Description Default num_samples Number of samples to draw 1 Returns: Type Description Samples Source code in normflows/distributions/target.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 def sample ( self , num_samples = 1 ): \"\"\"Sample from image distribution through rejection sampling Args: num_samples: Number of samples to draw Returns: Samples \"\"\" z = torch . zeros ( ( 0 , self . n_dims ), dtype = self . prop_scale . dtype , device = self . prop_scale . device ) while len ( z ) < num_samples : z_ = self . rejection_sampling ( num_samples ) ind = np . min ([ len ( z_ ), num_samples - len ( z )]) z = torch . cat ([ z , z_ [: ind , :]], 0 ) return z","title":"sample()"},{"location":"references/#normflows.distributions.target.TwoIndependent","text":"Bases: Target Target distribution that combines two independent distributions of equal size into one distribution. This is needed for Augmented Normalizing Flows, see https://arxiv.org/abs/2002.07101 Source code in normflows/distributions/target.py 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 class TwoIndependent ( Target ): \"\"\" Target distribution that combines two independent distributions of equal size into one distribution. This is needed for Augmented Normalizing Flows, see https://arxiv.org/abs/2002.07101 \"\"\" def __init__ ( self , target1 , target2 ): super () . __init__ () self . target1 = target1 self . target2 = target2 self . split = Split ( mode = 'channel' ) def log_prob ( self , z ): z1 , z2 = self . split ( z )[ 0 ] return self . target1 . log_prob ( z1 ) + self . target2 . log_prob ( z2 ) def sample ( self , num_samples = 1 ): z1 = self . target1 . sample ( num_samples ) z2 = self . target2 . sample ( num_samples ) return self . split . inverse ([ z1 , z2 ])[ 0 ]","title":"TwoIndependent"},{"location":"references/#normflows.distributions.target.TwoMoons","text":"Bases: Target Bimodal two-dimensional distribution Source code in normflows/distributions/target.py 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 class TwoMoons ( Target ): \"\"\" Bimodal two-dimensional distribution \"\"\" def __init__ ( self ): super () . __init__ () self . n_dims = 2 self . max_log_prob = 0.0 def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2 + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2) + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - 2 ) / 0.2 ) ** 2 - 0.5 * (( a - 2 ) / 0.3 ) ** 2 + torch . log ( 1 + torch . exp ( - 4 * a / 0.09 )) ) return log_prob","title":"TwoMoons"},{"location":"references/#normflows.distributions.target.TwoMoons.log_prob","text":"log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2 + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2) + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2)) Parameters: Name Type Description Default z value or batch of latent variable required Returns: Type Description log probability of the distribution for z Source code in normflows/distributions/target.py 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 def log_prob ( self , z ): \"\"\" ``` log(p) = - 1/2 * ((norm(z) - 2) / 0.2) ** 2 + log( exp(-1/2 * ((z[0] - 2) / 0.3) ** 2) + exp(-1/2 * ((z[0] + 2) / 0.3) ** 2)) ``` Args: z: value or batch of latent variable Returns: log probability of the distribution for z \"\"\" a = torch . abs ( z [:, 0 ]) log_prob = ( - 0.5 * (( torch . norm ( z , dim = 1 ) - 2 ) / 0.2 ) ** 2 - 0.5 * (( a - 2 ) / 0.3 ) ** 2 + torch . log ( 1 + torch . exp ( - 4 * a / 0.09 )) ) return log_prob","title":"log_prob()"},{"location":"references/#normflows.distributions.target_test","text":"","title":"target_test"},{"location":"references/#normflows.flows","text":"","title":"flows"},{"location":"references/#normflows.flows.affine","text":"","title":"affine"},{"location":"references/#normflows.flows.affine.autoregressive","text":"","title":"autoregressive"},{"location":"references/#normflows.flows.affine.autoregressive.Autoregressive","text":"Bases: Flow Transforms each input variable with an invertible elementwise transformation. The parameters of each invertible elementwise transformation can be functions of previous input variables, but they must not depend on the current or any following input variables. NOTE Calculating the inverse transform is D times slower than calculating the forward transform, where D is the dimensionality of the input to the transform. Source code in normflows/flows/affine/autoregressive.py 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class Autoregressive ( Flow ): \"\"\"Transforms each input variable with an invertible elementwise transformation. The parameters of each invertible elementwise transformation can be functions of previous input variables, but they must not depend on the current or any following input variables. **NOTE** Calculating the inverse transform is D times slower than calculating the forward transform, where D is the dimensionality of the input to the transform. \"\"\" def __init__ ( self , autoregressive_net ): super ( Autoregressive , self ) . __init__ () self . autoregressive_net = autoregressive_net def forward ( self , inputs , context = None ): autoregressive_params = self . autoregressive_net ( inputs , context ) outputs , logabsdet = self . _elementwise_forward ( inputs , autoregressive_params ) return outputs , logabsdet def inverse ( self , inputs , context = None ): num_inputs = np . prod ( inputs . shape [ 1 :]) outputs = torch . zeros_like ( inputs ) logabsdet = None for _ in range ( num_inputs ): autoregressive_params = self . autoregressive_net ( outputs , context ) outputs , logabsdet = self . _elementwise_inverse ( inputs , autoregressive_params ) return outputs , logabsdet def _output_dim_multiplier ( self ): raise NotImplementedError () def _elementwise_forward ( self , inputs , autoregressive_params ): raise NotImplementedError () def _elementwise_inverse ( self , inputs , autoregressive_params ): raise NotImplementedError ()","title":"Autoregressive"},{"location":"references/#normflows.flows.affine.autoregressive_test","text":"","title":"autoregressive_test"},{"location":"references/#normflows.flows.affine.coupling","text":"","title":"coupling"},{"location":"references/#normflows.flows.affine.coupling.AffineConstFlow","text":"Bases: Flow scales and shifts with learned constants per dimension. In the NICE paper there is a scaling layer which is a special case of this where t is None Source code in normflows/flows/affine/coupling.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 class AffineConstFlow ( Flow ): \"\"\" scales and shifts with learned constants per dimension. In the NICE paper there is a scaling layer which is a special case of this where t is None \"\"\" def __init__ ( self , shape , scale = True , shift = True ): \"\"\"Constructor Args: shape: Shape of the coupling layer scale: Flag whether to apply scaling shift: Flag whether to apply shift logscale_factor: Optional factor which can be used to control the scale of the log scale factor \"\"\" super () . __init__ () if scale : self . s = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"s\" , torch . zeros ( shape )[ None ]) if shift : self . t = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"t\" , torch . zeros ( shape )[ None ]) self . n_dim = self . s . dim () self . batch_dims = torch . nonzero ( torch . tensor ( self . s . shape ) == 1 , as_tuple = False )[:, 0 ] . tolist () def forward ( self , z ): z_ = z * torch . exp ( self . s ) + self . t if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = prod_batch_dims * torch . sum ( self . s ) return z_ , log_det def inverse ( self , z ): z_ = ( z - self . t ) * torch . exp ( - self . s ) if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = - prod_batch_dims * torch . sum ( self . s ) return z_ , log_det","title":"AffineConstFlow"},{"location":"references/#normflows.flows.affine.coupling.AffineConstFlow.__init__","text":"Constructor Parameters: Name Type Description Default shape Shape of the coupling layer required scale Flag whether to apply scaling True shift Flag whether to apply shift True logscale_factor Optional factor which can be used to control the scale of the log scale factor required Source code in normflows/flows/affine/coupling.py 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 def __init__ ( self , shape , scale = True , shift = True ): \"\"\"Constructor Args: shape: Shape of the coupling layer scale: Flag whether to apply scaling shift: Flag whether to apply shift logscale_factor: Optional factor which can be used to control the scale of the log scale factor \"\"\" super () . __init__ () if scale : self . s = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"s\" , torch . zeros ( shape )[ None ]) if shift : self . t = nn . Parameter ( torch . zeros ( shape )[ None ]) else : self . register_buffer ( \"t\" , torch . zeros ( shape )[ None ]) self . n_dim = self . s . dim () self . batch_dims = torch . nonzero ( torch . tensor ( self . s . shape ) == 1 , as_tuple = False )[:, 0 ] . tolist ()","title":"__init__()"},{"location":"references/#normflows.flows.affine.coupling.AffineCoupling","text":"Bases: Flow Affine Coupling layer as introduced RealNVP paper, see arXiv: 1605.08803 Source code in normflows/flows/affine/coupling.py 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 class AffineCoupling ( Flow ): \"\"\" Affine Coupling layer as introduced RealNVP paper, see arXiv: 1605.08803 \"\"\" def __init__ ( self , param_map , scale = True , scale_map = \"exp\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model \"\"\" super () . __init__ () self . add_module ( \"param_map\" , param_map ) self . scale = scale self . scale_map = scale_map def forward ( self , z ): \"\"\" z is a list of z1 and z2; ```z = [z1, z2]``` z1 is left constant and affine map is applied to z2 with parameters depending on z1 Args: z \"\"\" z1 , z2 = z param = self . param_map ( z1 ) if self . scale : shift = param [:, 0 :: 2 , ... ] scale_ = param [:, 1 :: 2 , ... ] if self . scale_map == \"exp\" : z2 = z2 * torch . exp ( scale_ ) + shift log_det = torch . sum ( scale_ , dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 / scale + shift log_det = - torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid_inv\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 * scale + shift log_det = torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) else : raise NotImplementedError ( \"This scale map is not implemented.\" ) else : z2 = z2 + param log_det = zero_log_det_like_z ( z2 ) return [ z1 , z2 ], log_det def inverse ( self , z ): z1 , z2 = z param = self . param_map ( z1 ) if self . scale : shift = param [:, 0 :: 2 , ... ] scale_ = param [:, 1 :: 2 , ... ] if self . scale_map == \"exp\" : z2 = ( z2 - shift ) * torch . exp ( - scale_ ) log_det = - torch . sum ( scale_ , dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = ( z2 - shift ) * scale log_det = torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid_inv\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = ( z2 - shift ) / scale log_det = - torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) else : raise NotImplementedError ( \"This scale map is not implemented.\" ) else : z2 = z2 - param log_det = zero_log_det_like_z ( z2 ) return [ z1 , z2 ], log_det","title":"AffineCoupling"},{"location":"references/#normflows.flows.affine.coupling.AffineCoupling.__init__","text":"Constructor Parameters: Name Type Description Default param_map Maps features to shift and scale parameter (if applicable) required scale Flag whether scale shall be applied True scale_map Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model 'exp' Source code in normflows/flows/affine/coupling.py 104 105 106 107 108 109 110 111 112 113 114 115 def __init__ ( self , param_map , scale = True , scale_map = \"exp\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow, 'sigmoid_inv' uses multiplicative sigmoid scale when sampling from the model \"\"\" super () . __init__ () self . add_module ( \"param_map\" , param_map ) self . scale = scale self . scale_map = scale_map","title":"__init__()"},{"location":"references/#normflows.flows.affine.coupling.AffineCoupling.forward","text":"z is a list of z1 and z2; z = [z1, z2] z1 is left constant and affine map is applied to z2 with parameters depending on z1 Source code in normflows/flows/affine/coupling.py 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 def forward ( self , z ): \"\"\" z is a list of z1 and z2; ```z = [z1, z2]``` z1 is left constant and affine map is applied to z2 with parameters depending on z1 Args: z \"\"\" z1 , z2 = z param = self . param_map ( z1 ) if self . scale : shift = param [:, 0 :: 2 , ... ] scale_ = param [:, 1 :: 2 , ... ] if self . scale_map == \"exp\" : z2 = z2 * torch . exp ( scale_ ) + shift log_det = torch . sum ( scale_ , dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 / scale + shift log_det = - torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) elif self . scale_map == \"sigmoid_inv\" : scale = torch . sigmoid ( scale_ + 2 ) z2 = z2 * scale + shift log_det = torch . sum ( torch . log ( scale ), dim = list ( range ( 1 , shift . dim ()))) else : raise NotImplementedError ( \"This scale map is not implemented.\" ) else : z2 = z2 + param log_det = zero_log_det_like_z ( z2 ) return [ z1 , z2 ], log_det","title":"forward()"},{"location":"references/#normflows.flows.affine.coupling.AffineCouplingBlock","text":"Bases: Flow Affine Coupling layer including split and merge operation Source code in normflows/flows/affine/coupling.py 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 class AffineCouplingBlock ( Flow ): \"\"\" Affine Coupling layer including split and merge operation \"\"\" def __init__ ( self , param_map , scale = True , scale_map = \"exp\" , split_mode = \"channel\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Split layer self . flows += [ Split ( split_mode )] # Affine coupling layer self . flows += [ AffineCoupling ( param_map , scale , scale_map )] # Merge layer self . flows += [ Merge ( split_mode )] def forward ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for flow in self . flows : z , log_det = flow ( z ) log_det_tot += log_det return z , log_det_tot def inverse ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_det_tot += log_det return z , log_det_tot","title":"AffineCouplingBlock"},{"location":"references/#normflows.flows.affine.coupling.AffineCouplingBlock.__init__","text":"Constructor Parameters: Name Type Description Default param_map Maps features to shift and scale parameter (if applicable) required scale Flag whether scale shall be applied True scale_map Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow 'exp' split_mode Splitting mode, for possible values see Split class 'channel' Source code in normflows/flows/affine/coupling.py 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 def __init__ ( self , param_map , scale = True , scale_map = \"exp\" , split_mode = \"channel\" ): \"\"\"Constructor Args: param_map: Maps features to shift and scale parameter (if applicable) scale: Flag whether scale shall be applied scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Split layer self . flows += [ Split ( split_mode )] # Affine coupling layer self . flows += [ AffineCoupling ( param_map , scale , scale_map )] # Merge layer self . flows += [ Merge ( split_mode )]","title":"__init__()"},{"location":"references/#normflows.flows.affine.coupling.CCAffineConst","text":"Bases: Flow Affine constant flow layer with class-conditional parameters Source code in normflows/flows/affine/coupling.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 class CCAffineConst ( Flow ): \"\"\" Affine constant flow layer with class-conditional parameters \"\"\" def __init__ ( self , shape , num_classes ): super () . __init__ () if isinstance ( shape , int ): shape = ( shape ,) self . shape = shape self . s = nn . Parameter ( torch . zeros ( shape )[ None ]) self . t = nn . Parameter ( torch . zeros ( shape )[ None ]) self . s_cc = nn . Parameter ( torch . zeros ( num_classes , np . prod ( shape ))) self . t_cc = nn . Parameter ( torch . zeros ( num_classes , np . prod ( shape ))) self . n_dim = self . s . dim () self . batch_dims = torch . nonzero ( torch . tensor ( self . s . shape ) == 1 , as_tuple = False )[:, 0 ] . tolist () def forward ( self , z , y ): s = self . s + ( y @ self . s_cc ) . view ( - 1 , * self . shape ) t = self . t + ( y @ self . t_cc ) . view ( - 1 , * self . shape ) z_ = z * torch . exp ( s ) + t if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = prod_batch_dims * torch . sum ( s , dim = list ( range ( 1 , self . n_dim ))) return z_ , log_det def inverse ( self , z , y ): s = self . s + ( y @ self . s_cc ) . view ( - 1 , * self . shape ) t = self . t + ( y @ self . t_cc ) . view ( - 1 , * self . shape ) z_ = ( z - t ) * torch . exp ( - s ) if len ( self . batch_dims ) > 1 : prod_batch_dims = np . prod ([ z . size ( i ) for i in self . batch_dims [ 1 :]]) else : prod_batch_dims = 1 log_det = - prod_batch_dims * torch . sum ( s , dim = list ( range ( 1 , self . n_dim ))) return z_ , log_det","title":"CCAffineConst"},{"location":"references/#normflows.flows.affine.coupling.MaskedAffineFlow","text":"Bases: Flow RealNVP as introduced in arXiv: 1605.08803 Masked affine flow: f(z) = b * z + (1 - b) * (z * exp(s(b * z)) + t) class AffineHalfFlow(Flow): is MaskedAffineFlow with alternating bit mask NICE is AffineFlow with only shifts (volume preserving) Source code in normflows/flows/affine/coupling.py 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 class MaskedAffineFlow ( Flow ): \"\"\"RealNVP as introduced in [arXiv: 1605.08803](https://arxiv.org/abs/1605.08803) Masked affine flow: ``` f(z) = b * z + (1 - b) * (z * exp(s(b * z)) + t) ``` - class AffineHalfFlow(Flow): is MaskedAffineFlow with alternating bit mask - NICE is AffineFlow with only shifts (volume preserving) \"\"\" def __init__ ( self , b , t = None , s = None ): \"\"\"Constructor Args: b: mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s t: translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied s: scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied \"\"\" super () . __init__ () self . b_cpu = b . view ( 1 , * b . size ()) self . register_buffer ( \"b\" , self . b_cpu ) if s is None : self . s = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"s\" , s ) if t is None : self . t = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"t\" , t ) def forward ( self , z ): z_masked = self . b * z scale = self . s ( z_masked ) nan = torch . tensor ( np . nan , dtype = z . dtype , device = z . device ) scale = torch . where ( torch . isfinite ( scale ), scale , nan ) trans = self . t ( z_masked ) trans = torch . where ( torch . isfinite ( trans ), trans , nan ) z_ = z_masked + ( 1 - self . b ) * ( z * torch . exp ( scale ) + trans ) log_det = torch . sum (( 1 - self . b ) * scale , dim = list ( range ( 1 , self . b . dim ()))) return z_ , log_det def inverse ( self , z ): z_masked = self . b * z scale = self . s ( z_masked ) nan = torch . tensor ( np . nan , dtype = z . dtype , device = z . device ) scale = torch . where ( torch . isfinite ( scale ), scale , nan ) trans = self . t ( z_masked ) trans = torch . where ( torch . isfinite ( trans ), trans , nan ) z_ = z_masked + ( 1 - self . b ) * ( z - trans ) * torch . exp ( - scale ) log_det = - torch . sum (( 1 - self . b ) * scale , dim = list ( range ( 1 , self . b . dim ()))) return z_ , log_det","title":"MaskedAffineFlow"},{"location":"references/#normflows.flows.affine.coupling.MaskedAffineFlow.__init__","text":"Constructor Parameters: Name Type Description Default b mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s required t translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied None s scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied None Source code in normflows/flows/affine/coupling.py 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 def __init__ ( self , b , t = None , s = None ): \"\"\"Constructor Args: b: mask for features, i.e. tensor of same size as latent data point filled with 0s and 1s t: translation mapping, i.e. neural network, where first input dimension is batch dim, if None no translation is applied s: scale mapping, i.e. neural network, where first input dimension is batch dim, if None no scale is applied \"\"\" super () . __init__ () self . b_cpu = b . view ( 1 , * b . size ()) self . register_buffer ( \"b\" , self . b_cpu ) if s is None : self . s = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"s\" , s ) if t is None : self . t = lambda x : torch . zeros_like ( x ) else : self . add_module ( \"t\" , t )","title":"__init__()"},{"location":"references/#normflows.flows.affine.coupling_test","text":"","title":"coupling_test"},{"location":"references/#normflows.flows.affine.glow","text":"","title":"glow"},{"location":"references/#normflows.flows.affine.glow.GlowBlock","text":"Bases: Flow Glow: Generative Flow with Invertible 1\u00d71 Convolutions, arXiv: 1807.03039 One Block of the Glow model, comprised of MaskedAffineFlow (affine coupling layer) Invertible1x1Conv (dropped if there is only one channel) ActNorm (first batch used for initialization) Source code in normflows/flows/affine/glow.py 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 class GlowBlock ( Flow ): \"\"\"Glow: Generative Flow with Invertible 1\u00d71 Convolutions, [arXiv: 1807.03039](https://arxiv.org/abs/1807.03039) One Block of the Glow model, comprised of - MaskedAffineFlow (affine coupling layer) - Invertible1x1Conv (dropped if there is only one channel) - ActNorm (first batch used for initialization) \"\"\" def __init__ ( self , channels , hidden_channels , scale = True , scale_map = \"sigmoid\" , split_mode = \"channel\" , leaky = 0.0 , init_zeros = True , use_lu = True , net_actnorm = False , ): \"\"\"Constructor Args: channels: Number of channels of the data hidden_channels: number of channels in the hidden layer of the ConvNet scale: Flag, whether to include scale in affine coupling layer scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class leaky: Leaky parameter of LeakyReLUs of ConvNet2d init_zeros: Flag whether to initialize last conv layer with zeros use_lu: Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers logscale_factor: Factor which can be used to control the scale of the log scale factor, see [source](https://github.com/openai/glow) \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Coupling layer kernel_size = ( 3 , 1 , 3 ) num_param = 2 if scale else 1 if \"channel\" == split_mode : channels_ = (( channels + 1 ) // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * ( channels // 2 ),) elif \"channel_inv\" == split_mode : channels_ = ( channels // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * (( channels + 1 ) // 2 ),) elif \"checkerboard\" in split_mode : channels_ = ( channels ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * channels ,) else : raise NotImplementedError ( \"Mode \" + split_mode + \" is not implemented.\" ) param_map = nets . ConvNet2d ( channels_ , kernel_size , leaky , init_zeros , actnorm = net_actnorm ) self . flows += [ AffineCouplingBlock ( param_map , scale , scale_map , split_mode )] # Invertible 1x1 convolution if channels > 1 : self . flows += [ Invertible1x1Conv ( channels , use_lu )] # Activation normalization self . flows += [ ActNorm (( channels ,) + ( 1 , 1 ))] def forward ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for flow in self . flows : z , log_det = flow ( z ) log_det_tot += log_det return z , log_det_tot def inverse ( self , z ): log_det_tot = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) for i in range ( len ( self . flows ) - 1 , - 1 , - 1 ): z , log_det = self . flows [ i ] . inverse ( z ) log_det_tot += log_det return z , log_det_tot","title":"GlowBlock"},{"location":"references/#normflows.flows.affine.glow.GlowBlock.__init__","text":"Constructor Parameters: Name Type Description Default channels Number of channels of the data required hidden_channels number of channels in the hidden layer of the ConvNet required scale Flag, whether to include scale in affine coupling layer True scale_map Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow 'sigmoid' split_mode Splitting mode, for possible values see Split class 'channel' leaky Leaky parameter of LeakyReLUs of ConvNet2d 0.0 init_zeros Flag whether to initialize last conv layer with zeros True use_lu Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers True logscale_factor Factor which can be used to control the scale of the log scale factor, see source required Source code in normflows/flows/affine/glow.py 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 def __init__ ( self , channels , hidden_channels , scale = True , scale_map = \"sigmoid\" , split_mode = \"channel\" , leaky = 0.0 , init_zeros = True , use_lu = True , net_actnorm = False , ): \"\"\"Constructor Args: channels: Number of channels of the data hidden_channels: number of channels in the hidden layer of the ConvNet scale: Flag, whether to include scale in affine coupling layer scale_map: Map to be applied to the scale parameter, can be 'exp' as in RealNVP or 'sigmoid' as in Glow split_mode: Splitting mode, for possible values see Split class leaky: Leaky parameter of LeakyReLUs of ConvNet2d init_zeros: Flag whether to initialize last conv layer with zeros use_lu: Flag whether to parametrize weights through the LU decomposition in invertible 1x1 convolution layers logscale_factor: Factor which can be used to control the scale of the log scale factor, see [source](https://github.com/openai/glow) \"\"\" super () . __init__ () self . flows = nn . ModuleList ([]) # Coupling layer kernel_size = ( 3 , 1 , 3 ) num_param = 2 if scale else 1 if \"channel\" == split_mode : channels_ = (( channels + 1 ) // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * ( channels // 2 ),) elif \"channel_inv\" == split_mode : channels_ = ( channels // 2 ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * (( channels + 1 ) // 2 ),) elif \"checkerboard\" in split_mode : channels_ = ( channels ,) + 2 * ( hidden_channels ,) channels_ += ( num_param * channels ,) else : raise NotImplementedError ( \"Mode \" + split_mode + \" is not implemented.\" ) param_map = nets . ConvNet2d ( channels_ , kernel_size , leaky , init_zeros , actnorm = net_actnorm ) self . flows += [ AffineCouplingBlock ( param_map , scale , scale_map , split_mode )] # Invertible 1x1 convolution if channels > 1 : self . flows += [ Invertible1x1Conv ( channels , use_lu )] # Activation normalization self . flows += [ ActNorm (( channels ,) + ( 1 , 1 ))]","title":"__init__()"},{"location":"references/#normflows.flows.affine.glow_test","text":"","title":"glow_test"},{"location":"references/#normflows.flows.base","text":"","title":"base"},{"location":"references/#normflows.flows.base.Composite","text":"Bases: Flow Composes several flows into one, in the order they are given. Source code in normflows/flows/base.py 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 class Composite ( Flow ): \"\"\" Composes several flows into one, in the order they are given. \"\"\" def __init__ ( self , flows ): \"\"\"Constructor Args: flows: Iterable of flows to composite \"\"\" super () . __init__ () self . _flows = nn . ModuleList ( flows ) @staticmethod def _cascade ( inputs , funcs ): batch_size = inputs . shape [ 0 ] outputs = inputs total_logabsdet = torch . zeros ( batch_size ) for func in funcs : outputs , logabsdet = func ( outputs ) total_logabsdet += logabsdet return outputs , total_logabsdet def forward ( self , inputs ): funcs = self . _flows return self . _cascade ( inputs , funcs ) def inverse ( self , inputs ): funcs = ( flow . inverse for flow in self . _flows [:: - 1 ]) return self . _cascade ( inputs , funcs )","title":"Composite"},{"location":"references/#normflows.flows.base.Composite.__init__","text":"Constructor Parameters: Name Type Description Default flows Iterable of flows to composite required Source code in normflows/flows/base.py 53 54 55 56 57 58 59 60 def __init__ ( self , flows ): \"\"\"Constructor Args: flows: Iterable of flows to composite \"\"\" super () . __init__ () self . _flows = nn . ModuleList ( flows )","title":"__init__()"},{"location":"references/#normflows.flows.base.Flow","text":"Bases: nn . Module Generic class for flow functions Source code in normflows/flows/base.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 class Flow ( nn . Module ): \"\"\" Generic class for flow functions \"\"\" def __init__ ( self ): super () . __init__ () def forward ( self , z ): \"\"\" Args: z: input variable, first dimension is batch dim Returns: transformed z and log of absolute determinant \"\"\" raise NotImplementedError ( \"Forward pass has not been implemented.\" ) def inverse ( self , z ): raise NotImplementedError ( \"This flow has no algebraic inverse.\" )","title":"Flow"},{"location":"references/#normflows.flows.base.Flow.forward","text":"Parameters: Name Type Description Default z input variable, first dimension is batch dim required Returns: Type Description transformed z and log of absolute determinant Source code in normflows/flows/base.py 13 14 15 16 17 18 19 20 21 def forward ( self , z ): \"\"\" Args: z: input variable, first dimension is batch dim Returns: transformed z and log of absolute determinant \"\"\" raise NotImplementedError ( \"Forward pass has not been implemented.\" )","title":"forward()"},{"location":"references/#normflows.flows.base.Reverse","text":"Bases: Flow Switches the forward transform of a flow layer with its inverse and vice versa Source code in normflows/flows/base.py 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 class Reverse ( Flow ): \"\"\" Switches the forward transform of a flow layer with its inverse and vice versa \"\"\" def __init__ ( self , flow ): \"\"\"Constructor Args: flow: Flow layer to be reversed \"\"\" super () . __init__ () self . flow = flow def forward ( self , z ): return self . flow . inverse ( z ) def inverse ( self , z ): return self . flow . forward ( z )","title":"Reverse"},{"location":"references/#normflows.flows.base.Reverse.__init__","text":"Constructor Parameters: Name Type Description Default flow Flow layer to be reversed required Source code in normflows/flows/base.py 32 33 34 35 36 37 38 39 def __init__ ( self , flow ): \"\"\"Constructor Args: flow: Flow layer to be reversed \"\"\" super () . __init__ () self . flow = flow","title":"__init__()"},{"location":"references/#normflows.flows.base_test","text":"","title":"base_test"},{"location":"references/#normflows.flows.flow_test","text":"","title":"flow_test"},{"location":"references/#normflows.flows.flow_test.FlowTest","text":"Bases: unittest . TestCase Generic test case for flow modules Source code in normflows/flows/flow_test.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 class FlowTest ( unittest . TestCase ): \"\"\" Generic test case for flow modules \"\"\" def assertClose ( self , actual , expected , atol = None , rtol = None ): assert_close ( actual , expected , atol = atol , rtol = rtol ) def checkForward ( self , flow , inputs ): # Do forward transform outputs , log_det = flow ( inputs ) # Check type assert outputs . dtype == inputs . dtype # Check shape assert outputs . shape == inputs . shape # Return results return outputs , log_det def checkInverse ( self , flow , inputs ): # Do inverse transform outputs , log_det = flow . inverse ( inputs ) # Check type assert outputs . dtype == inputs . dtype # Check shape assert outputs . shape == inputs . shape # Return results return outputs , log_det def checkForwardInverse ( self , flow , inputs , atol = None , rtol = None ): # Check forward outputs , log_det = self . checkForward ( flow , inputs ) # Check inverse input_ , log_det_ = self . checkInverse ( flow , outputs ) # Check identity self . assertClose ( input_ , inputs , atol , rtol ) ld_id = log_det + log_det_ self . assertClose ( ld_id , torch . zeros_like ( ld_id ), atol , rtol )","title":"FlowTest"},{"location":"references/#normflows.flows.mixing","text":"","title":"mixing"},{"location":"references/#normflows.flows.mixing.Invertible1x1Conv","text":"Bases: Flow Invertible 1x1 convolution introduced in the Glow paper Assumes 4d input/output tensors of the form NCHW Source code in normflows/flows/mixing.py 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 class Invertible1x1Conv ( Flow ): \"\"\" Invertible 1x1 convolution introduced in the Glow paper Assumes 4d input/output tensors of the form NCHW \"\"\" def __init__ ( self , num_channels , use_lu = False ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q ) def _assemble_W ( self , inverse = False ): # assemble W from its components (P, L, U, S) L = torch . tril ( self . L , diagonal =- 1 ) + self . eye U = torch . triu ( self . U , diagonal = 1 ) + torch . diag ( self . sign_S * torch . exp ( self . log_S ) ) if inverse : if self . log_S . dtype == torch . float64 : L_inv = torch . inverse ( L ) U_inv = torch . inverse ( U ) else : L_inv = torch . inverse ( L . double ()) . type ( self . log_S . dtype ) U_inv = torch . inverse ( U . double ()) . type ( self . log_S . dtype ) W = U_inv @ L_inv @ self . P . t () else : W = self . P @ L @ U return W def forward ( self , z ): if self . use_lu : W = self . _assemble_W ( inverse = True ) log_det = - torch . sum ( self . log_S ) else : W_dtype = self . W . dtype if W_dtype == torch . float64 : W = torch . inverse ( self . W ) else : W = torch . inverse ( self . W . double ()) . type ( W_dtype ) W = W . view ( * W . size (), 1 , 1 ) log_det = - torch . slogdet ( self . W )[ 1 ] W = W . view ( self . num_channels , self . num_channels , 1 , 1 ) z_ = torch . nn . functional . conv2d ( z , W ) log_det = log_det * z . size ( 2 ) * z . size ( 3 ) return z_ , log_det def inverse ( self , z ): if self . use_lu : W = self . _assemble_W () log_det = torch . sum ( self . log_S ) else : W = self . W log_det = torch . slogdet ( self . W )[ 1 ] W = W . view ( self . num_channels , self . num_channels , 1 , 1 ) z_ = torch . nn . functional . conv2d ( z , W ) log_det = log_det * z . size ( 2 ) * z . size ( 3 ) return z_ , log_det","title":"Invertible1x1Conv"},{"location":"references/#normflows.flows.mixing.Invertible1x1Conv.__init__","text":"Constructor Parameters: Name Type Description Default num_channels Number of channels of the data required use_lu Flag whether to parametrize weights through the LU decomposition False Source code in normflows/flows/mixing.py 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 def __init__ ( self , num_channels , use_lu = False ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q )","title":"__init__()"},{"location":"references/#normflows.flows.mixing.InvertibleAffine","text":"Bases: Flow Invertible affine transformation without shift, i.e. one-dimensional version of the invertible 1x1 convolutions Source code in normflows/flows/mixing.py 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 class InvertibleAffine ( Flow ): \"\"\" Invertible affine transformation without shift, i.e. one-dimensional version of the invertible 1x1 convolutions \"\"\" def __init__ ( self , num_channels , use_lu = True ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q ) def _assemble_W ( self , inverse = False ): # assemble W from its components (P, L, U, S) L = torch . tril ( self . L , diagonal =- 1 ) + self . eye U = torch . triu ( self . U , diagonal = 1 ) + torch . diag ( self . sign_S * torch . exp ( self . log_S ) ) if inverse : if self . log_S . dtype == torch . float64 : L_inv = torch . inverse ( L ) U_inv = torch . inverse ( U ) else : L_inv = torch . inverse ( L . double ()) . type ( self . log_S . dtype ) U_inv = torch . inverse ( U . double ()) . type ( self . log_S . dtype ) W = U_inv @ L_inv @ self . P . t () else : W = self . P @ L @ U return W def forward ( self , z ): if self . use_lu : W = self . _assemble_W ( inverse = True ) log_det = - torch . sum ( self . log_S ) else : W_dtype = self . W . dtype if W_dtype == torch . float64 : W = torch . inverse ( self . W ) else : W = torch . inverse ( self . W . double ()) . type ( W_dtype ) log_det = - torch . slogdet ( self . W )[ 1 ] z_ = z @ W return z_ , log_det def inverse ( self , z ): if self . use_lu : W = self . _assemble_W () log_det = torch . sum ( self . log_S ) else : W = self . W log_det = torch . slogdet ( self . W )[ 1 ] z_ = z @ W return z_ , log_det","title":"InvertibleAffine"},{"location":"references/#normflows.flows.mixing.InvertibleAffine.__init__","text":"Constructor Parameters: Name Type Description Default num_channels Number of channels of the data required use_lu Flag whether to parametrize weights through the LU decomposition True Source code in normflows/flows/mixing.py 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 def __init__ ( self , num_channels , use_lu = True ): \"\"\"Constructor Args: num_channels: Number of channels of the data use_lu: Flag whether to parametrize weights through the LU decomposition \"\"\" super () . __init__ () self . num_channels = num_channels self . use_lu = use_lu Q , _ = torch . linalg . qr ( torch . randn ( self . num_channels , self . num_channels )) if use_lu : P , L , U = torch . lu_unpack ( * Q . lu ()) self . register_buffer ( \"P\" , P ) # remains fixed during optimization self . L = nn . Parameter ( L ) # lower triangular portion S = U . diag () # \"crop out\" the diagonal to its own parameter self . register_buffer ( \"sign_S\" , torch . sign ( S )) self . log_S = nn . Parameter ( torch . log ( torch . abs ( S ))) self . U = nn . Parameter ( torch . triu ( U , diagonal = 1 ) ) # \"crop out\" diagonal, stored in S self . register_buffer ( \"eye\" , torch . diag ( torch . ones ( self . num_channels ))) else : self . W = nn . Parameter ( Q )","title":"__init__()"},{"location":"references/#normflows.flows.mixing.LULinearPermute","text":"Bases: Flow Fixed permutation combined with a linear transformation parametrized using the LU decomposition, used in https://arxiv.org/abs/1906.04032 Source code in normflows/flows/mixing.py 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 class LULinearPermute ( Flow ): \"\"\" Fixed permutation combined with a linear transformation parametrized using the LU decomposition, used in https://arxiv.org/abs/1906.04032 \"\"\" def __init__ ( self , num_channels , identity_init = True ): \"\"\"Constructor Args: num_channels: Number of dimensions of the data identity_init: Flag, whether to initialize linear transform as identity matrix \"\"\" # Initialize super () . __init__ () # Define modules self . permutation = _RandomPermutation ( num_channels ) self . linear = _LULinear ( num_channels , identity_init = identity_init ) def forward ( self , z ): z , log_det = self . linear . inverse ( z ) z , _ = self . permutation . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , _ = self . permutation ( z ) z , log_det = self . linear ( z ) return z , log_det . view ( - 1 )","title":"LULinearPermute"},{"location":"references/#normflows.flows.mixing.LULinearPermute.__init__","text":"Constructor Parameters: Name Type Description Default num_channels Number of dimensions of the data required identity_init Flag, whether to initialize linear transform as identity matrix True Source code in normflows/flows/mixing.py 541 542 543 544 545 546 547 548 549 550 551 552 553 def __init__ ( self , num_channels , identity_init = True ): \"\"\"Constructor Args: num_channels: Number of dimensions of the data identity_init: Flag, whether to initialize linear transform as identity matrix \"\"\" # Initialize super () . __init__ () # Define modules self . permutation = _RandomPermutation ( num_channels ) self . linear = _LULinear ( num_channels , identity_init = identity_init )","title":"__init__()"},{"location":"references/#normflows.flows.mixing.Permute","text":"Bases: Flow Permutation features along the channel dimension Source code in normflows/flows/mixing.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 class Permute ( Flow ): \"\"\" Permutation features along the channel dimension \"\"\" def __init__ ( self , num_channels , mode = \"shuffle\" ): \"\"\"Constructor Args: num_channel: Number of channels mode: Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part \"\"\" super () . __init__ () self . mode = mode self . num_channels = num_channels if self . mode == \"shuffle\" : perm = torch . randperm ( self . num_channels ) inv_perm = torch . empty_like ( perm ) . scatter_ ( dim = 0 , index = perm , src = torch . arange ( self . num_channels ) ) self . register_buffer ( \"perm\" , perm ) self . register_buffer ( \"inv_perm\" , inv_perm ) def forward ( self , z ): if self . mode == \"shuffle\" : z = z [:, self . perm , ... ] elif self . mode == \"swap\" : z1 = z [:, : self . num_channels // 2 , ... ] z2 = z [:, self . num_channels // 2 :, ... ] z = torch . cat ([ z2 , z1 ], dim = 1 ) else : raise NotImplementedError ( \"The mode \" + self . mode + \" is not implemented.\" ) log_det = torch . zeros ( len ( z ), device = z . device ) return z , log_det def inverse ( self , z ): if self . mode == \"shuffle\" : z = z [:, self . inv_perm , ... ] elif self . mode == \"swap\" : z1 = z [:, : ( self . num_channels + 1 ) // 2 , ... ] z2 = z [:, ( self . num_channels + 1 ) // 2 :, ... ] z = torch . cat ([ z2 , z1 ], dim = 1 ) else : raise NotImplementedError ( \"The mode \" + self . mode + \" is not implemented.\" ) log_det = torch . zeros ( len ( z ), device = z . device ) return z , log_det","title":"Permute"},{"location":"references/#normflows.flows.mixing.Permute.__init__","text":"Constructor Parameters: Name Type Description Default num_channel Number of channels required mode Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part 'shuffle' Source code in normflows/flows/mixing.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 def __init__ ( self , num_channels , mode = \"shuffle\" ): \"\"\"Constructor Args: num_channel: Number of channels mode: Mode of permuting features, can be shuffle for random permutation or swap for interchanging upper and lower part \"\"\" super () . __init__ () self . mode = mode self . num_channels = num_channels if self . mode == \"shuffle\" : perm = torch . randperm ( self . num_channels ) inv_perm = torch . empty_like ( perm ) . scatter_ ( dim = 0 , index = perm , src = torch . arange ( self . num_channels ) ) self . register_buffer ( \"perm\" , perm ) self . register_buffer ( \"inv_perm\" , inv_perm )","title":"__init__()"},{"location":"references/#normflows.flows.mixing_test","text":"","title":"mixing_test"},{"location":"references/#normflows.flows.neural_spline","text":"","title":"neural_spline"},{"location":"references/#normflows.flows.neural_spline.autoregressive","text":"Implementations of autoregressive transforms. Code taken from https://github.com/bayesiains/nsf","title":"autoregressive"},{"location":"references/#normflows.flows.neural_spline.autoregressive_test","text":"Tests for the autoregressive transforms. Code partially taken from https://github.com/bayesiains/nsf","title":"autoregressive_test"},{"location":"references/#normflows.flows.neural_spline.coupling","text":"Implementations of various coupling layers. Code taken from https://github.com/bayesiains/nsf","title":"coupling"},{"location":"references/#normflows.flows.neural_spline.coupling.Coupling","text":"Bases: Flow A base class for coupling layers. Supports 2D inputs (NxD), as well as 4D inputs for images (NxCxHxW). For images the splitting is done on the channel dimension, using the provided 1D mask. Source code in normflows/flows/neural_spline/coupling.py 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 class Coupling ( Flow ): \"\"\"A base class for coupling layers. Supports 2D inputs (NxD), as well as 4D inputs for images (NxCxHxW). For images the splitting is done on the channel dimension, using the provided 1D mask.\"\"\" def __init__ ( self , mask , transform_net_create_fn , unconditional_transform = None ): \"\"\"Constructor. mask: a 1-dim tensor, tuple or list. It indexes inputs as follows: - if `mask[i] > 0`, `input[i]` will be transformed. - if `mask[i] <= 0`, `input[i]` will be passed unchanged. Args: mask \"\"\" mask = torch . as_tensor ( mask ) if mask . dim () != 1 : raise ValueError ( \"Mask must be a 1-dim tensor.\" ) if mask . numel () <= 0 : raise ValueError ( \"Mask can't be empty.\" ) super () . __init__ () self . features = len ( mask ) features_vector = torch . arange ( self . features ) self . register_buffer ( \"identity_features\" , features_vector . masked_select ( mask <= 0 ) ) self . register_buffer ( \"transform_features\" , features_vector . masked_select ( mask > 0 ) ) assert self . num_identity_features + self . num_transform_features == self . features self . transform_net = transform_net_create_fn ( self . num_identity_features , self . num_transform_features * self . _transform_dim_multiplier (), ) if unconditional_transform is None : self . unconditional_transform = None else : self . unconditional_transform = unconditional_transform ( features = self . num_identity_features ) @property def num_identity_features ( self ): return len ( self . identity_features ) @property def num_transform_features ( self ): return len ( self . transform_features ) def forward ( self , inputs , context = None ): if inputs . dim () not in [ 2 , 4 ]: raise ValueError ( \"Inputs must be a 2D or a 4D tensor.\" ) if inputs . shape [ 1 ] != self . features : raise ValueError ( \"Expected features = {} , got {} .\" . format ( self . features , inputs . shape [ 1 ]) ) identity_split = inputs [:, self . identity_features , ... ] transform_split = inputs [:, self . transform_features , ... ] transform_params = self . transform_net ( identity_split , context ) transform_split , logabsdet = self . _coupling_transform_forward ( inputs = transform_split , transform_params = transform_params ) if self . unconditional_transform is not None : identity_split , logabsdet_identity = self . unconditional_transform ( identity_split , context ) logabsdet += logabsdet_identity outputs = torch . empty_like ( inputs ) outputs [:, self . identity_features , ... ] = identity_split outputs [:, self . transform_features , ... ] = transform_split return outputs , logabsdet def inverse ( self , inputs , context = None ): if inputs . dim () not in [ 2 , 4 ]: raise ValueError ( \"Inputs must be a 2D or a 4D tensor.\" ) if inputs . shape [ 1 ] != self . features : raise ValueError ( \"Expected features = {} , got {} .\" . format ( self . features , inputs . shape [ 1 ]) ) identity_split = inputs [:, self . identity_features , ... ] transform_split = inputs [:, self . transform_features , ... ] logabsdet = 0.0 if self . unconditional_transform is not None : identity_split , logabsdet = self . unconditional_transform . inverse ( identity_split , context ) transform_params = self . transform_net ( identity_split , context ) transform_split , logabsdet_split = self . _coupling_transform_inverse ( inputs = transform_split , transform_params = transform_params ) logabsdet += logabsdet_split outputs = torch . empty_like ( inputs ) outputs [:, self . identity_features ] = identity_split outputs [:, self . transform_features ] = transform_split return outputs , logabsdet def _transform_dim_multiplier ( self ): \"\"\"Number of features to output for each transform dimension.\"\"\" raise NotImplementedError () def _coupling_transform_forward ( self , inputs , transform_params ): \"\"\"Forward pass of the coupling transform.\"\"\" raise NotImplementedError () def _coupling_transform_inverse ( self , inputs , transform_params ): \"\"\"Inverse of the coupling transform.\"\"\" raise NotImplementedError ()","title":"Coupling"},{"location":"references/#normflows.flows.neural_spline.coupling.Coupling.__init__","text":"Constructor. mask: a 1-dim tensor, tuple or list. It indexes inputs as follows: if mask[i] > 0 , input[i] will be transformed. if mask[i] <= 0 , input[i] will be passed unchanged. Source code in normflows/flows/neural_spline/coupling.py 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 def __init__ ( self , mask , transform_net_create_fn , unconditional_transform = None ): \"\"\"Constructor. mask: a 1-dim tensor, tuple or list. It indexes inputs as follows: - if `mask[i] > 0`, `input[i]` will be transformed. - if `mask[i] <= 0`, `input[i]` will be passed unchanged. Args: mask \"\"\" mask = torch . as_tensor ( mask ) if mask . dim () != 1 : raise ValueError ( \"Mask must be a 1-dim tensor.\" ) if mask . numel () <= 0 : raise ValueError ( \"Mask can't be empty.\" ) super () . __init__ () self . features = len ( mask ) features_vector = torch . arange ( self . features ) self . register_buffer ( \"identity_features\" , features_vector . masked_select ( mask <= 0 ) ) self . register_buffer ( \"transform_features\" , features_vector . masked_select ( mask > 0 ) ) assert self . num_identity_features + self . num_transform_features == self . features self . transform_net = transform_net_create_fn ( self . num_identity_features , self . num_transform_features * self . _transform_dim_multiplier (), ) if unconditional_transform is None : self . unconditional_transform = None else : self . unconditional_transform = unconditional_transform ( features = self . num_identity_features )","title":"__init__()"},{"location":"references/#normflows.flows.neural_spline.coupling_test","text":"Tests for the coupling Transforms. Code partially taken from https://github.com/bayesiains/nsf","title":"coupling_test"},{"location":"references/#normflows.flows.neural_spline.wrapper","text":"","title":"wrapper"},{"location":"references/#normflows.flows.neural_spline.wrapper.AutoregressiveRationalQuadraticSpline","text":"Bases: Flow Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see sources Source code in normflows/flows/neural_spline/wrapper.py 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 class AutoregressiveRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see [sources](https://github.com/bayesiains/nsf) \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = False , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = \"linear\" , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , ) def forward ( self , z ): z , log_det = self . mprqat . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . mprqat ( z ) return z , log_det . view ( - 1 )","title":"AutoregressiveRationalQuadraticSpline"},{"location":"references/#normflows.flows.neural_spline.wrapper.AutoregressiveRationalQuadraticSpline.__init__","text":"Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required num_bins int Number of bins 8 tail_bound int Bound of the spline tails 3 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 permute_mask bool Flag, permutes the mask of the NN False init_identity bool Flag, initialize transform as identity True Source code in normflows/flows/neural_spline/wrapper.py 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = False , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = \"linear\" , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , )","title":"__init__()"},{"location":"references/#normflows.flows.neural_spline.wrapper.CircularAutoregressiveRationalQuadraticSpline","text":"Bases: Flow Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see sources Source code in normflows/flows/neural_spline/wrapper.py 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 class CircularAutoregressiveRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see [sources](https://github.com/bayesiains/nsf) \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = True , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = tails , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , ) def forward ( self , z ): z , log_det = self . mprqat . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . mprqat ( z ) return z , log_det . view ( - 1 )","title":"CircularAutoregressiveRationalQuadraticSpline"},{"location":"references/#normflows.flows.neural_spline.wrapper.CircularAutoregressiveRationalQuadraticSpline.__init__","text":"Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required ind_circ Iterable Indices of the circular coordinates required num_bins int Number of bins 8 tail_bound int Bound of the spline tails 3 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 permute_mask bool Flag, permutes the mask of the NN True init_identity bool Flag, initialize transform as identity True Source code in normflows/flows/neural_spline/wrapper.py 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3 , activation = nn . ReLU , dropout_probability = 0.0 , permute_mask = True , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (int): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN permute_mask (bool): Flag, permutes the mask of the NN init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . mprqat = MaskedPiecewiseRationalQuadraticAutoregressive ( features = num_input_channels , hidden_features = num_hidden_channels , context_features = None , num_bins = num_bins , tails = tails , tail_bound = tail_bound , num_blocks = num_blocks , use_residual_blocks = True , random_mask = False , permute_mask = permute_mask , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , init_identity = init_identity , )","title":"__init__()"},{"location":"references/#normflows.flows.neural_spline.wrapper.CircularCoupledRationalQuadraticSpline","text":"Bases: Flow Neural spline flow coupling layer with circular coordinates Source code in normflows/flows/neural_spline/wrapper.py 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 class CircularCoupledRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer with circular coordinates \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , mask = None , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (float or Iterable): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used mask (torch tensor): Mask to be used, alternating masked generated is None init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () if mask is None : mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ) features_vector = torch . arange ( num_input_channels ) identity_features = features_vector . masked_select ( mask <= 0 ) ind_circ = torch . tensor ( ind_circ ) ind_circ_id = [] for i , id in enumerate ( identity_features ): if id in ind_circ : ind_circ_id += [ i ] if torch . is_tensor ( tail_bound ): scale_pf = np . pi / tail_bound [ ind_circ_id ] else : scale_pf = np . pi / tail_bound def transform_net_create_fn ( in_features , out_features ): if len ( ind_circ_id ) > 0 : pf = PeriodicFeaturesElementwise ( in_features , ind_circ_id , scale_pf ) else : pf = None net = ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , preprocessing = pf , ) if init_identity : torch . nn . init . constant_ ( net . final_layer . weight , 0.0 ) torch . nn . init . constant_ ( net . final_layer . bias , np . log ( np . exp ( 1 - DEFAULT_MIN_DERIVATIVE ) - 1 ) ) return net tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . prqct = PiecewiseRationalQuadraticCoupling ( mask = mask , transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , apply_unconditional_transform = True , ) def forward ( self , z ): z , log_det = self . prqct . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . prqct ( z ) return z , log_det . view ( - 1 )","title":"CircularCoupledRationalQuadraticSpline"},{"location":"references/#normflows.flows.neural_spline.wrapper.CircularCoupledRationalQuadraticSpline.__init__","text":"Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required ind_circ Iterable Indices of the circular coordinates required num_bins int Number of bins 8 tail_bound float or Iterable Bound of the spline tails 3.0 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 reverse_mask bool Flag whether the reverse mask should be used False mask torch tensor Mask to be used, alternating masked generated is None None init_identity bool Flag, initialize transform as identity True Source code in normflows/flows/neural_spline/wrapper.py 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , ind_circ , num_bins = 8 , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , mask = None , init_identity = True , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN ind_circ (Iterable): Indices of the circular coordinates num_bins (int): Number of bins tail_bound (float or Iterable): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used mask (torch tensor): Mask to be used, alternating masked generated is None init_identity (bool): Flag, initialize transform as identity \"\"\" super () . __init__ () if mask is None : mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ) features_vector = torch . arange ( num_input_channels ) identity_features = features_vector . masked_select ( mask <= 0 ) ind_circ = torch . tensor ( ind_circ ) ind_circ_id = [] for i , id in enumerate ( identity_features ): if id in ind_circ : ind_circ_id += [ i ] if torch . is_tensor ( tail_bound ): scale_pf = np . pi / tail_bound [ ind_circ_id ] else : scale_pf = np . pi / tail_bound def transform_net_create_fn ( in_features , out_features ): if len ( ind_circ_id ) > 0 : pf = PeriodicFeaturesElementwise ( in_features , ind_circ_id , scale_pf ) else : pf = None net = ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , preprocessing = pf , ) if init_identity : torch . nn . init . constant_ ( net . final_layer . weight , 0.0 ) torch . nn . init . constant_ ( net . final_layer . bias , np . log ( np . exp ( 1 - DEFAULT_MIN_DERIVATIVE ) - 1 ) ) return net tails = [ \"circular\" if i in ind_circ else \"linear\" for i in range ( num_input_channels ) ] self . prqct = PiecewiseRationalQuadraticCoupling ( mask = mask , transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , apply_unconditional_transform = True , )","title":"__init__()"},{"location":"references/#normflows.flows.neural_spline.wrapper.CoupledRationalQuadraticSpline","text":"Bases: Flow Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see source Source code in normflows/flows/neural_spline/wrapper.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 class CoupledRationalQuadraticSpline ( Flow ): \"\"\" Neural spline flow coupling layer, wrapper for the implementation of Durkan et al., see [source](https://github.com/bayesiains/nsf) \"\"\" def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tails = \"linear\" , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tails (str): Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval tail_bound (float): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used \"\"\" super () . __init__ () def transform_net_create_fn ( in_features , out_features ): return ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , ) self . prqct = PiecewiseRationalQuadraticCoupling ( mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ), transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , # Setting True corresponds to equations (4), (5), (6) in the NSF paper: apply_unconditional_transform = True , ) def forward ( self , z ): z , log_det = self . prqct . inverse ( z ) return z , log_det . view ( - 1 ) def inverse ( self , z ): z , log_det = self . prqct ( z ) return z , log_det . view ( - 1 )","title":"CoupledRationalQuadraticSpline"},{"location":"references/#normflows.flows.neural_spline.wrapper.CoupledRationalQuadraticSpline.__init__","text":"Constructor Parameters: Name Type Description Default num_input_channels int Flow dimension required num_blocks int Number of residual blocks of the parameter NN required num_hidden_channels int Number of hidden units of the NN required num_bins int Number of bins 8 tails str Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval 'linear' tail_bound float Bound of the spline tails 3.0 activation torch module Activation function nn.ReLU dropout_probability float Dropout probability of the NN 0.0 reverse_mask bool Flag whether the reverse mask should be used False Source code in normflows/flows/neural_spline/wrapper.py 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 def __init__ ( self , num_input_channels , num_blocks , num_hidden_channels , num_bins = 8 , tails = \"linear\" , tail_bound = 3.0 , activation = nn . ReLU , dropout_probability = 0.0 , reverse_mask = False , ): \"\"\"Constructor Args: num_input_channels (int): Flow dimension num_blocks (int): Number of residual blocks of the parameter NN num_hidden_channels (int): Number of hidden units of the NN num_bins (int): Number of bins tails (str): Behaviour of the tails of the distribution, can be linear, circular for periodic distribution, or None for distribution on the compact interval tail_bound (float): Bound of the spline tails activation (torch module): Activation function dropout_probability (float): Dropout probability of the NN reverse_mask (bool): Flag whether the reverse mask should be used \"\"\" super () . __init__ () def transform_net_create_fn ( in_features , out_features ): return ResidualNet ( in_features = in_features , out_features = out_features , context_features = None , hidden_features = num_hidden_channels , num_blocks = num_blocks , activation = activation (), dropout_probability = dropout_probability , use_batch_norm = False , ) self . prqct = PiecewiseRationalQuadraticCoupling ( mask = create_alternating_binary_mask ( num_input_channels , even = reverse_mask ), transform_net_create_fn = transform_net_create_fn , num_bins = num_bins , tails = tails , tail_bound = tail_bound , # Setting True corresponds to equations (4), (5), (6) in the NSF paper: apply_unconditional_transform = True , )","title":"__init__()"},{"location":"references/#normflows.flows.neural_spline.wrapper_test","text":"","title":"wrapper_test"},{"location":"references/#normflows.flows.normalization","text":"","title":"normalization"},{"location":"references/#normflows.flows.normalization.ActNorm","text":"Bases: AffineConstFlow An AffineConstFlow but with a data-dependent initialization, where on the very first batch we clever initialize the s,t so that the output is unit gaussian. As described in Glow paper. Source code in normflows/flows/normalization.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 class ActNorm ( AffineConstFlow ): \"\"\" An AffineConstFlow but with a data-dependent initialization, where on the very first batch we clever initialize the s,t so that the output is unit gaussian. As described in Glow paper. \"\"\" def __init__ ( self , * args , ** kwargs ): super () . __init__ ( * args , ** kwargs ) self . data_dep_init_done_cpu = torch . tensor ( 0.0 ) self . register_buffer ( \"data_dep_init_done\" , self . data_dep_init_done_cpu ) def forward ( self , z ): # first batch is used for initialization, c.f. batchnorm if not self . data_dep_init_done > 0.0 : assert self . s is not None and self . t is not None s_init = - torch . log ( z . std ( dim = self . batch_dims , keepdim = True ) + 1e-6 ) self . s . data = s_init . data self . t . data = ( - z . mean ( dim = self . batch_dims , keepdim = True ) * torch . exp ( self . s ) ) . data self . data_dep_init_done = torch . tensor ( 1.0 ) return super () . forward ( z ) def inverse ( self , z ): # first batch is used for initialization, c.f. batchnorm if not self . data_dep_init_done : assert self . s is not None and self . t is not None s_init = torch . log ( z . std ( dim = self . batch_dims , keepdim = True ) + 1e-6 ) self . s . data = s_init . data self . t . data = z . mean ( dim = self . batch_dims , keepdim = True ) . data self . data_dep_init_done = torch . tensor ( 1.0 ) return super () . inverse ( z )","title":"ActNorm"},{"location":"references/#normflows.flows.normalization.BatchNorm","text":"Bases: Flow Batch Normalization with out considering the derivatives of the batch statistics, see arXiv: 1605.08803 Source code in normflows/flows/normalization.py 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 class BatchNorm ( Flow ): \"\"\" Batch Normalization with out considering the derivatives of the batch statistics, see [arXiv: 1605.08803](https://arxiv.org/abs/1605.08803) \"\"\" def __init__ ( self , eps = 1.0e-10 ): super () . __init__ () self . eps_cpu = torch . tensor ( eps ) self . register_buffer ( \"eps\" , self . eps_cpu ) def forward ( self , z ): \"\"\" Do batch norm over batch and sample dimension \"\"\" mean = torch . mean ( z , dim = 0 , keepdims = True ) std = torch . std ( z , dim = 0 , keepdims = True ) z_ = ( z - mean ) / torch . sqrt ( std ** 2 + self . eps ) log_det = torch . log ( 1 / torch . prod ( torch . sqrt ( std ** 2 + self . eps ))) . repeat ( z . size ()[ 0 ] ) return z_ , log_det","title":"BatchNorm"},{"location":"references/#normflows.flows.normalization.BatchNorm.forward","text":"Do batch norm over batch and sample dimension Source code in normflows/flows/normalization.py 52 53 54 55 56 57 58 59 60 61 62 def forward ( self , z ): \"\"\" Do batch norm over batch and sample dimension \"\"\" mean = torch . mean ( z , dim = 0 , keepdims = True ) std = torch . std ( z , dim = 0 , keepdims = True ) z_ = ( z - mean ) / torch . sqrt ( std ** 2 + self . eps ) log_det = torch . log ( 1 / torch . prod ( torch . sqrt ( std ** 2 + self . eps ))) . repeat ( z . size ()[ 0 ] ) return z_ , log_det","title":"forward()"},{"location":"references/#normflows.flows.periodic","text":"","title":"periodic"},{"location":"references/#normflows.flows.periodic.PeriodicShift","text":"Bases: Flow Shift and wrap periodic coordinates Source code in normflows/flows/periodic.py 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 class PeriodicShift ( Flow ): \"\"\" Shift and wrap periodic coordinates \"\"\" def __init__ ( self , ind , bound = 1.0 , shift = 0.0 ): \"\"\"Constructor Args: ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval shift: Tensor, shift to be applied \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound if torch . is_tensor ( shift ): self . register_buffer ( \"shift\" , shift ) else : self . shift = shift def forward ( self , z ): z_ = z . clone () z_ [ ... , self . ind ] = ( torch . remainder ( z_ [ ... , self . ind ] + self . shift + self . bound , 2 * self . bound ) - self . bound ) return z_ , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device ) def inverse ( self , z ): z_ = z . clone () z_ [ ... , self . ind ] = ( torch . remainder ( z_ [ ... , self . ind ] - self . shift + self . bound , 2 * self . bound ) - self . bound ) return z_ , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device )","title":"PeriodicShift"},{"location":"references/#normflows.flows.periodic.PeriodicShift.__init__","text":"Constructor Parameters: Name Type Description Default ind Iterable, indices of coordinates to be mapped required bound Float or iterable, bound of interval 1.0 shift Tensor, shift to be applied 0.0 Source code in normflows/flows/periodic.py 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 def __init__ ( self , ind , bound = 1.0 , shift = 0.0 ): \"\"\"Constructor Args: ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval shift: Tensor, shift to be applied \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound if torch . is_tensor ( shift ): self . register_buffer ( \"shift\" , shift ) else : self . shift = shift","title":"__init__()"},{"location":"references/#normflows.flows.periodic.PeriodicWrap","text":"Bases: Flow Map periodic coordinates to fixed interval Source code in normflows/flows/periodic.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 class PeriodicWrap ( Flow ): \"\"\" Map periodic coordinates to fixed interval \"\"\" def __init__ ( self , ind , bound = 1.0 ): \"\"\"Constructor ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound def forward ( self , z ): return z , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device ) def inverse ( self , z ): z_ = z . clone () z_ [ ... , self . ind ] = ( torch . remainder ( z_ [ ... , self . ind ] + self . bound , 2 * self . bound ) - self . bound ) return z_ , torch . zeros ( len ( z ), dtype = z . dtype , device = z . device )","title":"PeriodicWrap"},{"location":"references/#normflows.flows.periodic.PeriodicWrap.__init__","text":"Constructor ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval Source code in normflows/flows/periodic.py 11 12 13 14 15 16 17 18 19 20 21 22 def __init__ ( self , ind , bound = 1.0 ): \"\"\"Constructor ind: Iterable, indices of coordinates to be mapped bound: Float or iterable, bound of interval \"\"\" super () . __init__ () self . ind = ind if torch . is_tensor ( bound ): self . register_buffer ( \"bound\" , bound ) else : self . bound = bound","title":"__init__()"},{"location":"references/#normflows.flows.periodic_test","text":"","title":"periodic_test"},{"location":"references/#normflows.flows.planar","text":"","title":"planar"},{"location":"references/#normflows.flows.planar.Planar","text":"Bases: Flow Planar flow as introduced in arXiv: 1505.05770 f(z) = z + u * h(w * z + b) Source code in normflows/flows/planar.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 class Planar ( Flow ): \"\"\"Planar flow as introduced in [arXiv: 1505.05770](https://arxiv.org/abs/1505.05770) ``` f(z) = z + u * h(w * z + b) ``` \"\"\" def __init__ ( self , shape , act = \"tanh\" , u = None , w = None , b = None ): \"\"\"Constructor of the planar flow Args: shape: shape of the latent variable z h: nonlinear function h of the planar flow (see definition of f above) u,w,b: optional initialization for parameters \"\"\" super () . __init__ () lim_w = np . sqrt ( 2.0 / np . prod ( shape )) lim_u = np . sqrt ( 2 ) if u is not None : self . u = nn . Parameter ( u ) else : self . u = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . u , - lim_u , lim_u ) if w is not None : self . w = nn . Parameter ( w ) else : self . w = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . w , - lim_w , lim_w ) if b is not None : self . b = nn . Parameter ( b ) else : self . b = nn . Parameter ( torch . zeros ( 1 )) self . act = act if act == \"tanh\" : self . h = torch . tanh elif act == \"leaky_relu\" : self . h = torch . nn . LeakyReLU ( negative_slope = 0.2 ) else : raise NotImplementedError ( \"Nonlinearity is not implemented.\" ) def forward ( self , z ): lin = torch . sum ( self . w * z , list ( range ( 1 , self . w . dim ())), keepdim = True ) + self . b inner = torch . sum ( self . w * self . u ) u = self . u + ( torch . log ( 1 + torch . exp ( inner )) - 1 - inner ) \\ * self . w / torch . sum ( self . w ** 2 ) # constraint w.T * u > -1 if self . act == \"tanh\" : h_ = lambda x : 1 / torch . cosh ( x ) ** 2 elif self . act == \"leaky_relu\" : h_ = lambda x : ( x < 0 ) * ( self . h . negative_slope - 1.0 ) + 1.0 z_ = z + u * self . h ( lin ) log_det = torch . log ( torch . abs ( 1 + torch . sum ( self . w * u ) * h_ ( lin . reshape ( - 1 )))) return z_ , log_det def inverse ( self , z ): if self . act != \"leaky_relu\" : raise NotImplementedError ( \"This flow has no algebraic inverse.\" ) lin = torch . sum ( self . w * z , list ( range ( 1 , self . w . dim ()))) + self . b a = ( lin < 0 ) * ( self . h . negative_slope - 1.0 ) + 1.0 # absorb leakyReLU slope into u inner = torch . sum ( self . w * self . u ) u = self . u + ( torch . log ( 1 + torch . exp ( inner )) - 1 - inner ) \\ * self . w / torch . sum ( self . w ** 2 ) dims = [ - 1 ] + ( u . dim () - 1 ) * [ 1 ] u = a . reshape ( * dims ) * u inner_ = torch . sum ( self . w * u , list ( range ( 1 , self . w . dim ()))) z_ = z - u * ( lin / ( 1 + inner_ )) . reshape ( * dims ) log_det = - torch . log ( torch . abs ( 1 + inner_ )) return z_ , log_det","title":"Planar"},{"location":"references/#normflows.flows.planar.Planar.__init__","text":"Constructor of the planar flow Parameters: Name Type Description Default shape shape of the latent variable z required h nonlinear function h of the planar flow (see definition of f above) required u,w,b optional initialization for parameters required Source code in normflows/flows/planar.py 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 def __init__ ( self , shape , act = \"tanh\" , u = None , w = None , b = None ): \"\"\"Constructor of the planar flow Args: shape: shape of the latent variable z h: nonlinear function h of the planar flow (see definition of f above) u,w,b: optional initialization for parameters \"\"\" super () . __init__ () lim_w = np . sqrt ( 2.0 / np . prod ( shape )) lim_u = np . sqrt ( 2 ) if u is not None : self . u = nn . Parameter ( u ) else : self . u = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . u , - lim_u , lim_u ) if w is not None : self . w = nn . Parameter ( w ) else : self . w = nn . Parameter ( torch . empty ( shape )[ None ]) nn . init . uniform_ ( self . w , - lim_w , lim_w ) if b is not None : self . b = nn . Parameter ( b ) else : self . b = nn . Parameter ( torch . zeros ( 1 )) self . act = act if act == \"tanh\" : self . h = torch . tanh elif act == \"leaky_relu\" : self . h = torch . nn . LeakyReLU ( negative_slope = 0.2 ) else : raise NotImplementedError ( \"Nonlinearity is not implemented.\" )","title":"__init__()"},{"location":"references/#normflows.flows.planar_test","text":"","title":"planar_test"},{"location":"references/#normflows.flows.radial","text":"","title":"radial"},{"location":"references/#normflows.flows.radial.Radial","text":"Bases: Flow Radial flow as introduced in arXiv: 1505.05770 f(z) = z + beta * h(alpha, r) * (z - z_0) Source code in normflows/flows/radial.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 class Radial ( Flow ): \"\"\"Radial flow as introduced in [arXiv: 1505.05770](https://arxiv.org/abs/1505.05770) ``` f(z) = z + beta * h(alpha, r) * (z - z_0) ``` \"\"\" def __init__ ( self , shape , z_0 = None ): \"\"\"Constructor of the radial flow Args: shape: shape of the latent variable z z_0: parameter of the radial flow \"\"\" super () . __init__ () self . d_cpu = torch . prod ( torch . tensor ( shape )) self . register_buffer ( \"d\" , self . d_cpu ) self . beta = nn . Parameter ( torch . empty ( 1 )) lim = 1.0 / np . prod ( shape ) nn . init . uniform_ ( self . beta , - lim - 1.0 , lim - 1.0 ) self . alpha = nn . Parameter ( torch . empty ( 1 )) nn . init . uniform_ ( self . alpha , - lim , lim ) if z_0 is not None : self . z_0 = nn . Parameter ( z_0 ) else : self . z_0 = nn . Parameter ( torch . randn ( shape )[ None ]) def forward ( self , z ): beta = torch . log ( 1 + torch . exp ( self . beta )) - torch . abs ( self . alpha ) dz = z - self . z_0 r = torch . linalg . vector_norm ( dz , dim = list ( range ( 1 , self . z_0 . dim ())), keepdim = True ) h_arr = beta / ( torch . abs ( self . alpha ) + r ) h_arr_ = - beta * r / ( torch . abs ( self . alpha ) + r ) ** 2 z_ = z + h_arr * dz log_det = ( self . d - 1 ) * torch . log ( 1 + h_arr ) + torch . log ( 1 + h_arr + h_arr_ ) log_det = log_det . reshape ( - 1 ) return z_ , log_det","title":"Radial"},{"location":"references/#normflows.flows.radial.Radial.__init__","text":"Constructor of the radial flow Parameters: Name Type Description Default shape shape of the latent variable z required z_0 parameter of the radial flow None Source code in normflows/flows/radial.py 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 def __init__ ( self , shape , z_0 = None ): \"\"\"Constructor of the radial flow Args: shape: shape of the latent variable z z_0: parameter of the radial flow \"\"\" super () . __init__ () self . d_cpu = torch . prod ( torch . tensor ( shape )) self . register_buffer ( \"d\" , self . d_cpu ) self . beta = nn . Parameter ( torch . empty ( 1 )) lim = 1.0 / np . prod ( shape ) nn . init . uniform_ ( self . beta , - lim - 1.0 , lim - 1.0 ) self . alpha = nn . Parameter ( torch . empty ( 1 )) nn . init . uniform_ ( self . alpha , - lim , lim ) if z_0 is not None : self . z_0 = nn . Parameter ( z_0 ) else : self . z_0 = nn . Parameter ( torch . randn ( shape )[ None ])","title":"__init__()"},{"location":"references/#normflows.flows.radial_test","text":"","title":"radial_test"},{"location":"references/#normflows.flows.reshape","text":"","title":"reshape"},{"location":"references/#normflows.flows.reshape.Merge","text":"Bases: Split Same as Split but with forward and backward pass interchanged Source code in normflows/flows/reshape.py 88 89 90 91 92 93 94 95 96 97 98 99 100 class Merge ( Split ): \"\"\" Same as Split but with forward and backward pass interchanged \"\"\" def __init__ ( self , mode = \"channel\" ): super () . __init__ ( mode ) def forward ( self , z ): return super () . inverse ( z ) def inverse ( self , z ): return super () . forward ( z )","title":"Merge"},{"location":"references/#normflows.flows.reshape.Split","text":"Bases: Flow Split features into two sets Source code in normflows/flows/reshape.py 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 class Split ( Flow ): \"\"\" Split features into two sets \"\"\" def __init__ ( self , mode = \"channel\" ): \"\"\"Constructor The splitting mode can be: - channel: Splits first feature dimension, usually channels, into two halfs - channel_inv: Same as channel, but with z1 and z2 flipped - checkerboard: Splits features using a checkerboard pattern (last feature dimension must be even) - checkerboard_inv: Same as checkerboard, but with inverted coloring Args: mode: splitting mode \"\"\" super () . __init__ () self . mode = mode def forward ( self , z ): if self . mode == \"channel\" : z1 , z2 = z . chunk ( 2 , dim = 1 ) elif self . mode == \"channel_inv\" : z2 , z1 = z . chunk ( 2 , dim = 1 ) elif \"checkerboard\" in self . mode : n_dims = z . dim () cb0 = 0 cb1 = 1 for i in range ( 1 , n_dims ): cb0_ = cb0 cb1_ = cb1 cb0 = [ cb0_ if j % 2 == 0 else cb1_ for j in range ( z . size ( n_dims - i ))] cb1 = [ cb1_ if j % 2 == 0 else cb0_ for j in range ( z . size ( n_dims - i ))] cb = cb1 if \"inv\" in self . mode else cb0 cb = torch . tensor ( cb )[ None ] . repeat ( len ( z ), * (( n_dims - 1 ) * [ 1 ])) cb = cb . to ( z . device ) z_size = z . size () z1 = z . reshape ( - 1 )[ torch . nonzero ( cb . view ( - 1 ), as_tuple = False )] . view ( * z_size [: - 1 ], - 1 ) z2 = z . reshape ( - 1 )[ torch . nonzero (( 1 - cb ) . view ( - 1 ), as_tuple = False )] . view ( * z_size [: - 1 ], - 1 ) else : raise NotImplementedError ( \"Mode \" + self . mode + \" is not implemented.\" ) log_det = 0 return [ z1 , z2 ], log_det def inverse ( self , z ): z1 , z2 = z if self . mode == \"channel\" : z = torch . cat ([ z1 , z2 ], 1 ) elif self . mode == \"channel_inv\" : z = torch . cat ([ z2 , z1 ], 1 ) elif \"checkerboard\" in self . mode : n_dims = z1 . dim () z_size = list ( z1 . size ()) z_size [ - 1 ] *= 2 cb0 = 0 cb1 = 1 for i in range ( 1 , n_dims ): cb0_ = cb0 cb1_ = cb1 cb0 = [ cb0_ if j % 2 == 0 else cb1_ for j in range ( z_size [ n_dims - i ])] cb1 = [ cb1_ if j % 2 == 0 else cb0_ for j in range ( z_size [ n_dims - i ])] cb = cb1 if \"inv\" in self . mode else cb0 cb = torch . tensor ( cb )[ None ] . repeat ( z_size [ 0 ], * (( n_dims - 1 ) * [ 1 ])) cb = cb . to ( z1 . device ) z1 = z1 [ ... , None ] . repeat ( * ( n_dims * [ 1 ]), 2 ) . view ( * z_size [: - 1 ], - 1 ) z2 = z2 [ ... , None ] . repeat ( * ( n_dims * [ 1 ]), 2 ) . view ( * z_size [: - 1 ], - 1 ) z = cb * z1 + ( 1 - cb ) * z2 else : raise NotImplementedError ( \"Mode \" + self . mode + \" is not implemented.\" ) log_det = 0 return z , log_det","title":"Split"},{"location":"references/#normflows.flows.reshape.Split.__init__","text":"Constructor The splitting mode can be: channel: Splits first feature dimension, usually channels, into two halfs channel_inv: Same as channel, but with z1 and z2 flipped checkerboard: Splits features using a checkerboard pattern (last feature dimension must be even) checkerboard_inv: Same as checkerboard, but with inverted coloring Parameters: Name Type Description Default mode splitting mode 'channel' Source code in normflows/flows/reshape.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 def __init__ ( self , mode = \"channel\" ): \"\"\"Constructor The splitting mode can be: - channel: Splits first feature dimension, usually channels, into two halfs - channel_inv: Same as channel, but with z1 and z2 flipped - checkerboard: Splits features using a checkerboard pattern (last feature dimension must be even) - checkerboard_inv: Same as checkerboard, but with inverted coloring Args: mode: splitting mode \"\"\" super () . __init__ () self . mode = mode","title":"__init__()"},{"location":"references/#normflows.flows.reshape.Squeeze","text":"Bases: Flow Squeeze operation of multi-scale architecture, RealNVP or Glow paper Source code in normflows/flows/reshape.py 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 class Squeeze ( Flow ): \"\"\" Squeeze operation of multi-scale architecture, RealNVP or Glow paper \"\"\" def __init__ ( self ): \"\"\" Constructor \"\"\" super () . __init__ () def forward ( self , z ): log_det = 0 s = z . size () z = z . view ( s [ 0 ], s [ 1 ] // 4 , 2 , 2 , s [ 2 ], s [ 3 ]) z = z . permute ( 0 , 1 , 4 , 2 , 5 , 3 ) . contiguous () z = z . view ( s [ 0 ], s [ 1 ] // 4 , 2 * s [ 2 ], 2 * s [ 3 ]) return z , log_det def inverse ( self , z ): log_det = 0 s = z . size () z = z . view ( * s [: 2 ], s [ 2 ] // 2 , 2 , s [ 3 ] // 2 , 2 ) z = z . permute ( 0 , 1 , 3 , 5 , 2 , 4 ) . contiguous () z = z . view ( s [ 0 ], 4 * s [ 1 ], s [ 2 ] // 2 , s [ 3 ] // 2 ) return z , log_det","title":"Squeeze"},{"location":"references/#normflows.flows.reshape.Squeeze.__init__","text":"Constructor Source code in normflows/flows/reshape.py 108 109 110 111 112 def __init__ ( self ): \"\"\" Constructor \"\"\" super () . __init__ ()","title":"__init__()"},{"location":"references/#normflows.flows.residual","text":"","title":"residual"},{"location":"references/#normflows.flows.residual.Residual","text":"Bases: Flow Invertible residual net block, wrapper to the implementation of Chen et al., see sources Source code in normflows/flows/residual.py 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 class Residual ( Flow ): \"\"\" Invertible residual net block, wrapper to the implementation of Chen et al., see [sources](https://github.com/rtqichen/residual-flows) \"\"\" def __init__ ( self , net , reverse = True , reduce_memory = True , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" ): \"\"\"Constructor Args: net: Neural network, must be Lipschitz continuous with L < 1 reverse: Flag, if true the map ```f(x) = x + net(x)``` is applied in the inverse pass, otherwise it is done in forward reduce_memory: Flag, if true Neumann series and precomputations, for backward pass in forward pass are done geom_p: Parameter of the geometric distribution used for the Neumann series lamb: Parameter of the geometric distribution used for the Neumann series n_power_series: Number of terms in the Neumann series exact_trace: Flag, if true the trace of the Jacobian is computed exactly brute_force: Flag, if true the Jacobian is computed exactly in 2D n_samples: Number of samples used to estimate power series n_exact_terms: Number of terms always included in the power series n_dist: Distribution used for the power series, either \"geometric\" or \"poisson\" \"\"\" super () . __init__ () self . reverse = reverse self . iresblock = iResBlock ( net , n_samples = n_samples , n_exact_terms = n_exact_terms , neumann_grad = reduce_memory , grad_in_forward = reduce_memory , exact_trace = exact_trace , geom_p = geom_p , lamb = lamb , n_power_series = n_power_series , brute_force = brute_force , n_dist = n_dist , ) def forward ( self , z ): if self . reverse : z , log_det = self . iresblock . inverse ( z , 0 ) else : z , log_det = self . iresblock . forward ( z , 0 ) return z , - log_det . view ( - 1 ) def inverse ( self , z ): if self . reverse : z , log_det = self . iresblock . forward ( z , 0 ) else : z , log_det = self . iresblock . inverse ( z , 0 ) return z , - log_det . view ( - 1 )","title":"Residual"},{"location":"references/#normflows.flows.residual.Residual.__init__","text":"Constructor Parameters: Name Type Description Default net Neural network, must be Lipschitz continuous with L < 1 required reverse Flag, if true the map f(x) = x + net(x) is applied in the inverse pass, otherwise it is done in forward True reduce_memory Flag, if true Neumann series and precomputations, for backward pass in forward pass are done True geom_p Parameter of the geometric distribution used for the Neumann series 0.5 lamb Parameter of the geometric distribution used for the Neumann series 2.0 n_power_series Number of terms in the Neumann series None exact_trace Flag, if true the trace of the Jacobian is computed exactly False brute_force Flag, if true the Jacobian is computed exactly in 2D False n_samples Number of samples used to estimate power series 1 n_exact_terms Number of terms always included in the power series 2 n_dist Distribution used for the power series, either \"geometric\" or \"poisson\" 'geometric' Source code in normflows/flows/residual.py 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 def __init__ ( self , net , reverse = True , reduce_memory = True , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" ): \"\"\"Constructor Args: net: Neural network, must be Lipschitz continuous with L < 1 reverse: Flag, if true the map ```f(x) = x + net(x)``` is applied in the inverse pass, otherwise it is done in forward reduce_memory: Flag, if true Neumann series and precomputations, for backward pass in forward pass are done geom_p: Parameter of the geometric distribution used for the Neumann series lamb: Parameter of the geometric distribution used for the Neumann series n_power_series: Number of terms in the Neumann series exact_trace: Flag, if true the trace of the Jacobian is computed exactly brute_force: Flag, if true the Jacobian is computed exactly in 2D n_samples: Number of samples used to estimate power series n_exact_terms: Number of terms always included in the power series n_dist: Distribution used for the power series, either \"geometric\" or \"poisson\" \"\"\" super () . __init__ () self . reverse = reverse self . iresblock = iResBlock ( net , n_samples = n_samples , n_exact_terms = n_exact_terms , neumann_grad = reduce_memory , grad_in_forward = reduce_memory , exact_trace = exact_trace , geom_p = geom_p , lamb = lamb , n_power_series = n_power_series , brute_force = brute_force , n_dist = n_dist , )","title":"__init__()"},{"location":"references/#normflows.flows.residual.iResBlock","text":"Bases: nn . Module Source code in normflows/flows/residual.py 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 class iResBlock ( nn . Module ): def __init__ ( self , nnet , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" , neumann_grad = True , grad_in_forward = False , ): \"\"\" Args: nnet: a nn.Module n_power_series: number of power series. If not None, uses a biased approximation to logdet. exact_trace: if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. brute_force: Computes the exact logdet. Only available for 2D inputs. \"\"\" nn . Module . __init__ ( self ) self . nnet = nnet self . n_dist = n_dist self . geom_p = nn . Parameter ( torch . tensor ( np . log ( geom_p ) - np . log ( 1.0 - geom_p ))) self . lamb = nn . Parameter ( torch . tensor ( lamb )) self . n_samples = n_samples self . n_power_series = n_power_series self . exact_trace = exact_trace self . brute_force = brute_force self . n_exact_terms = n_exact_terms self . grad_in_forward = grad_in_forward self . neumann_grad = neumann_grad # store the samples of n. self . register_buffer ( \"last_n_samples\" , torch . zeros ( self . n_samples )) self . register_buffer ( \"last_firmom\" , torch . zeros ( 1 )) self . register_buffer ( \"last_secmom\" , torch . zeros ( 1 )) def forward ( self , x , logpx = None ): if logpx is None : y = x + self . nnet ( x ) return y else : g , logdetgrad = self . _logdetgrad ( x ) return x + g , logpx - logdetgrad def inverse ( self , y , logpy = None ): x = self . _inverse_fixed_point ( y ) if logpy is None : return x else : return x , logpy + self . _logdetgrad ( x )[ 1 ] def _inverse_fixed_point ( self , y , atol = 1e-5 , rtol = 1e-5 ): x , x_prev = y - self . nnet ( y ), y i = 0 tol = atol + y . abs () * rtol while not torch . all (( x - x_prev ) ** 2 / tol < 1 ): x , x_prev = y - self . nnet ( x ), x i += 1 if i > 1000 : break return x def _logdetgrad ( self , x ): \"\"\"Returns g(x) and ```logdet|d(x+g(x))/dx|```\"\"\" with torch . enable_grad (): if ( self . brute_force or not self . training ) and ( x . ndimension () == 2 and x . shape [ 1 ] == 2 ): ########################################### # Brute-force compute Jacobian determinant. ########################################### x = x . requires_grad_ ( True ) g = self . nnet ( x ) # Brute-force logdet only available for 2D. jac = batch_jacobian ( g , x ) batch_dets = ( jac [:, 0 , 0 ] + 1 ) * ( jac [:, 1 , 1 ] + 1 ) - jac [ :, 0 , 1 ] * jac [:, 1 , 0 ] return g , torch . log ( torch . abs ( batch_dets )) . view ( - 1 , 1 ) if self . n_dist == \"geometric\" : geom_p = torch . sigmoid ( self . geom_p ) . item () sample_fn = lambda m : geometric_sample ( geom_p , m ) rcdf_fn = lambda k , offset : geometric_1mcdf ( geom_p , k , offset ) elif self . n_dist == \"poisson\" : lamb = self . lamb . item () sample_fn = lambda m : poisson_sample ( lamb , m ) rcdf_fn = lambda k , offset : poisson_1mcdf ( lamb , k , offset ) if self . training : if self . n_power_series is None : # Unbiased estimation. lamb = self . lamb . item () n_samples = sample_fn ( self . n_samples ) n_power_series = max ( n_samples ) + self . n_exact_terms coeff_fn = ( lambda k : 1 / rcdf_fn ( k , self . n_exact_terms ) * sum ( n_samples >= k - self . n_exact_terms ) / len ( n_samples ) ) else : # Truncated estimation. n_power_series = self . n_power_series coeff_fn = lambda k : 1.0 else : # Unbiased estimation with more exact terms. lamb = self . lamb . item () n_samples = sample_fn ( self . n_samples ) n_power_series = max ( n_samples ) + 20 coeff_fn = ( lambda k : 1 / rcdf_fn ( k , 20 ) * sum ( n_samples >= k - 20 ) / len ( n_samples ) ) if not self . exact_trace : #################################### # Power series with trace estimator. #################################### vareps = torch . randn_like ( x ) # Choose the type of estimator. if self . training and self . neumann_grad : estimator_fn = neumann_logdet_estimator else : estimator_fn = basic_logdet_estimator # Do backprop-in-forward to save memory. if self . training and self . grad_in_forward : g , logdetgrad = mem_eff_wrapper ( estimator_fn , self . nnet , x , n_power_series , vareps , coeff_fn , self . training , ) else : x = x . requires_grad_ ( True ) g = self . nnet ( x ) logdetgrad = estimator_fn ( g , x , n_power_series , vareps , coeff_fn , self . training ) else : ############################################ # Power series with exact trace computation. ############################################ x = x . requires_grad_ ( True ) g = self . nnet ( x ) jac = batch_jacobian ( g , x ) logdetgrad = batch_trace ( jac ) jac_k = jac for k in range ( 2 , n_power_series + 1 ): jac_k = torch . bmm ( jac , jac_k ) logdetgrad = logdetgrad + ( - 1 ) ** ( k + 1 ) / k * coeff_fn ( k ) * batch_trace ( jac_k ) if self . training and self . n_power_series is None : self . last_n_samples . copy_ ( torch . tensor ( n_samples ) . to ( self . last_n_samples ) ) estimator = logdetgrad . detach () self . last_firmom . copy_ ( torch . mean ( estimator ) . to ( self . last_firmom )) self . last_secmom . copy_ ( torch . mean ( estimator ** 2 ) . to ( self . last_secmom )) return g , logdetgrad . view ( - 1 , 1 ) def extra_repr ( self ): return \"dist= {} , n_samples= {} , n_power_series= {} , neumann_grad= {} , exact_trace= {} , brute_force= {} \" . format ( self . n_dist , self . n_samples , self . n_power_series , self . neumann_grad , self . exact_trace , self . brute_force , )","title":"iResBlock"},{"location":"references/#normflows.flows.residual.iResBlock.__init__","text":"Parameters: Name Type Description Default nnet a nn.Module required n_power_series number of power series. If not None, uses a biased approximation to logdet. None exact_trace if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. False brute_force Computes the exact logdet. Only available for 2D inputs. False Source code in normflows/flows/residual.py 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 def __init__ ( self , nnet , geom_p = 0.5 , lamb = 2.0 , n_power_series = None , exact_trace = False , brute_force = False , n_samples = 1 , n_exact_terms = 2 , n_dist = \"geometric\" , neumann_grad = True , grad_in_forward = False , ): \"\"\" Args: nnet: a nn.Module n_power_series: number of power series. If not None, uses a biased approximation to logdet. exact_trace: if False, uses a Hutchinson trace estimator. Otherwise computes the exact full Jacobian. brute_force: Computes the exact logdet. Only available for 2D inputs. \"\"\" nn . Module . __init__ ( self ) self . nnet = nnet self . n_dist = n_dist self . geom_p = nn . Parameter ( torch . tensor ( np . log ( geom_p ) - np . log ( 1.0 - geom_p ))) self . lamb = nn . Parameter ( torch . tensor ( lamb )) self . n_samples = n_samples self . n_power_series = n_power_series self . exact_trace = exact_trace self . brute_force = brute_force self . n_exact_terms = n_exact_terms self . grad_in_forward = grad_in_forward self . neumann_grad = neumann_grad # store the samples of n. self . register_buffer ( \"last_n_samples\" , torch . zeros ( self . n_samples )) self . register_buffer ( \"last_firmom\" , torch . zeros ( 1 )) self . register_buffer ( \"last_secmom\" , torch . zeros ( 1 ))","title":"__init__()"},{"location":"references/#normflows.flows.residual_test","text":"","title":"residual_test"},{"location":"references/#normflows.flows.stochastic","text":"","title":"stochastic"},{"location":"references/#normflows.flows.stochastic.HamiltonianMonteCarlo","text":"Bases: Flow Flow layer using the HMC proposal in Stochastic Normalising Flows See arXiv: 2002.06707 Source code in normflows/flows/stochastic.py 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 class HamiltonianMonteCarlo ( Flow ): \"\"\"Flow layer using the HMC proposal in Stochastic Normalising Flows See [arXiv: 2002.06707](https://arxiv.org/abs/2002.06707) \"\"\" def __init__ ( self , target , steps , log_step_size , log_mass , max_abs_grad = None ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. steps: The number of leapfrog steps log_step_size: The log step size used in the leapfrog integrator. shape (dim) log_mass: The log_mass determining the variance of the momentum samples. shape (dim) max_abs_grad: Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability.\"\"\" super () . __init__ () self . target = target self . steps = steps self . register_parameter ( \"log_step_size\" , torch . nn . Parameter ( log_step_size )) self . register_parameter ( \"log_mass\" , torch . nn . Parameter ( log_mass )) self . max_abs_grad = max_abs_grad def forward ( self , z ): # Draw momentum p = torch . randn_like ( z ) * torch . exp ( 0.5 * self . log_mass ) # leapfrog z_new = z . clone () p_new = p . clone () step_size = torch . exp ( self . log_step_size ) for i in range ( self . steps ): p_half = p_new - ( step_size / 2.0 ) * - self . gradlogP ( z_new ) z_new = z_new + step_size * ( p_half / torch . exp ( self . log_mass )) p_new = p_half - ( step_size / 2.0 ) * - self . gradlogP ( z_new ) # Metropolis Hastings correction probabilities = torch . exp ( self . target . log_prob ( z_new ) - self . target . log_prob ( z ) - 0.5 * torch . sum ( p_new ** 2 / torch . exp ( self . log_mass ), 1 ) + 0.5 * torch . sum ( p ** 2 / torch . exp ( self . log_mass ), 1 ) ) uniforms = torch . rand_like ( probabilities ) mask = uniforms < probabilities z_out = torch . where ( mask . unsqueeze ( 1 ), z_new , z ) return z_out , self . target . log_prob ( z ) - self . target . log_prob ( z_out ) def inverse ( self , z ): return self . forward ( z ) def gradlogP ( self , z ): z_ = z . detach () . requires_grad_ () logp = self . target . log_prob ( z_ ) grad = torch . autograd . grad ( logp , z_ , grad_outputs = torch . ones_like ( logp ))[ 0 ] if self . max_abs_grad : grad = torch . clamp ( grad , max = self . max_abs_grad , min =- self . max_abs_grad ) return grad","title":"HamiltonianMonteCarlo"},{"location":"references/#normflows.flows.stochastic.HamiltonianMonteCarlo.__init__","text":"Constructor Parameters: Name Type Description Default target The stationary distribution of this Markov transition, i.e. the target distribution to sample from. required steps The number of leapfrog steps required log_step_size The log step size used in the leapfrog integrator. shape (dim) required log_mass The log_mass determining the variance of the momentum samples. shape (dim) required max_abs_grad Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability. None Source code in normflows/flows/stochastic.py 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 def __init__ ( self , target , steps , log_step_size , log_mass , max_abs_grad = None ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. steps: The number of leapfrog steps log_step_size: The log step size used in the leapfrog integrator. shape (dim) log_mass: The log_mass determining the variance of the momentum samples. shape (dim) max_abs_grad: Maximum absolute value of the gradient of the target distribution's log probability. If set to None then no gradient clipping is applied. Useful for improving numerical stability.\"\"\" super () . __init__ () self . target = target self . steps = steps self . register_parameter ( \"log_step_size\" , torch . nn . Parameter ( log_step_size )) self . register_parameter ( \"log_mass\" , torch . nn . Parameter ( log_mass )) self . max_abs_grad = max_abs_grad","title":"__init__()"},{"location":"references/#normflows.flows.stochastic.MetropolisHastings","text":"Bases: Flow Sampling through Metropolis Hastings in Stochastic Normalizing Flow See arXiv: 2002.06707 Source code in normflows/flows/stochastic.py 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class MetropolisHastings ( Flow ): \"\"\"Sampling through Metropolis Hastings in Stochastic Normalizing Flow See [arXiv: 2002.06707](https://arxiv.org/abs/2002.06707) \"\"\" def __init__ ( self , target , proposal , steps ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. proposal: Proposal distribution steps: Number of MCMC steps to perform \"\"\" super () . __init__ () self . target = target self . proposal = proposal self . steps = steps def forward ( self , z ): # Initialize number of samples and log(det) num_samples = len ( z ) log_det = torch . zeros ( num_samples , dtype = z . dtype , device = z . device ) # Get log(p) for current samples log_p = self . target . log_prob ( z ) for i in range ( self . steps ): # Make proposal and get log(p) z_ , log_p_diff = self . proposal ( z ) log_p_ = self . target . log_prob ( z_ ) # Make acceptance decision w = torch . rand ( num_samples , dtype = z . dtype , device = z . device ) log_w_accept = log_p_ - log_p + log_p_diff w_accept = torch . clamp ( torch . exp ( log_w_accept ), max = 1 ) accept = w <= w_accept # Update samples, log(det), and log(p) z = torch . where ( accept . unsqueeze ( 1 ), z_ , z ) log_det_ = log_p - log_p_ log_det = torch . where ( accept , log_det + log_det_ , log_det ) log_p = torch . where ( accept , log_p_ , log_p ) return z , log_det def inverse ( self , z ): # Equivalent to forward pass return self . forward ( z )","title":"MetropolisHastings"},{"location":"references/#normflows.flows.stochastic.MetropolisHastings.__init__","text":"Constructor Parameters: Name Type Description Default target The stationary distribution of this Markov transition, i.e. the target distribution to sample from. required proposal Proposal distribution required steps Number of MCMC steps to perform required Source code in normflows/flows/stochastic.py 12 13 14 15 16 17 18 19 20 21 22 23 def __init__ ( self , target , proposal , steps ): \"\"\"Constructor Args: target: The stationary distribution of this Markov transition, i.e. the target distribution to sample from. proposal: Proposal distribution steps: Number of MCMC steps to perform \"\"\" super () . __init__ () self . target = target self . proposal = proposal self . steps = steps","title":"__init__()"},{"location":"references/#normflows.flows.stochastic_test","text":"","title":"stochastic_test"},{"location":"references/#normflows.nets","text":"","title":"nets"},{"location":"references/#normflows.nets.cnn","text":"","title":"cnn"},{"location":"references/#normflows.nets.cnn.ConvNet2d","text":"Bases: nn . Module Convolutional Neural Network with leaky ReLU nonlinearities Source code in normflows/nets/cnn.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 class ConvNet2d ( nn . Module ): \"\"\" Convolutional Neural Network with leaky ReLU nonlinearities \"\"\" def __init__ ( self , channels , kernel_size , leaky = 0.0 , init_zeros = True , actnorm = False , weight_std = None , ): \"\"\"Constructor Args: channels: List of channels of conv layers, first entry is in_channels kernel_size: List of kernel sizes, same for height and width leaky: Leaky part of ReLU init_zeros: Flag whether last layer shall be initialized with zeros scale_output: Flag whether to scale output with a log scale parameter logscale_factor: Constant factor to be multiplied to log scaling actnorm: Flag whether activation normalization shall be done after each conv layer except output weight_std: Fixed std used to initialize every layer \"\"\" super () . __init__ () # Build network net = nn . ModuleList ([]) for i in range ( len ( kernel_size ) - 1 ): conv = nn . Conv2d ( channels [ i ], channels [ i + 1 ], kernel_size [ i ], padding = kernel_size [ i ] // 2 , bias = ( not actnorm ), ) if weight_std is not None : conv . weight . data . normal_ ( mean = 0.0 , std = weight_std ) net . append ( conv ) if actnorm : net . append ( utils . ActNorm (( channels [ i + 1 ],) + ( 1 , 1 ))) net . append ( nn . LeakyReLU ( leaky )) i = len ( kernel_size ) net . append ( nn . Conv2d ( channels [ i - 1 ], channels [ i ], kernel_size [ i - 1 ], padding = kernel_size [ i - 1 ] // 2 , ) ) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) self . net = nn . Sequential ( * net ) def forward ( self , x ): return self . net ( x )","title":"ConvNet2d"},{"location":"references/#normflows.nets.cnn.ConvNet2d.__init__","text":"Constructor Parameters: Name Type Description Default channels List of channels of conv layers, first entry is in_channels required kernel_size List of kernel sizes, same for height and width required leaky Leaky part of ReLU 0.0 init_zeros Flag whether last layer shall be initialized with zeros True scale_output Flag whether to scale output with a log scale parameter required logscale_factor Constant factor to be multiplied to log scaling required actnorm Flag whether activation normalization shall be done after each conv layer except output False weight_std Fixed std used to initialize every layer None Source code in normflows/nets/cnn.py 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 def __init__ ( self , channels , kernel_size , leaky = 0.0 , init_zeros = True , actnorm = False , weight_std = None , ): \"\"\"Constructor Args: channels: List of channels of conv layers, first entry is in_channels kernel_size: List of kernel sizes, same for height and width leaky: Leaky part of ReLU init_zeros: Flag whether last layer shall be initialized with zeros scale_output: Flag whether to scale output with a log scale parameter logscale_factor: Constant factor to be multiplied to log scaling actnorm: Flag whether activation normalization shall be done after each conv layer except output weight_std: Fixed std used to initialize every layer \"\"\" super () . __init__ () # Build network net = nn . ModuleList ([]) for i in range ( len ( kernel_size ) - 1 ): conv = nn . Conv2d ( channels [ i ], channels [ i + 1 ], kernel_size [ i ], padding = kernel_size [ i ] // 2 , bias = ( not actnorm ), ) if weight_std is not None : conv . weight . data . normal_ ( mean = 0.0 , std = weight_std ) net . append ( conv ) if actnorm : net . append ( utils . ActNorm (( channels [ i + 1 ],) + ( 1 , 1 ))) net . append ( nn . LeakyReLU ( leaky )) i = len ( kernel_size ) net . append ( nn . Conv2d ( channels [ i - 1 ], channels [ i ], kernel_size [ i - 1 ], padding = kernel_size [ i - 1 ] // 2 , ) ) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) self . net = nn . Sequential ( * net )","title":"__init__()"},{"location":"references/#normflows.nets.lipschitz","text":"","title":"lipschitz"},{"location":"references/#normflows.nets.lipschitz.LipschitzCNN","text":"Bases: nn . Module Convolutional neural network which is Lipschitz continuous with Lipschitz constant L < 1 Source code in normflows/nets/lipschitz.py 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 class LipschitzCNN ( nn . Module ): \"\"\" Convolutional neural network which is Lipschitz continuous with Lipschitz constant L < 1 \"\"\" def __init__ ( self , channels , kernel_size , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\"Constructor Args: channels: Integer list with the number of channels of the layers kernel_size: Integer list of kernel sizes of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( kernel_size ) self . channels = channels self . kernel_size = kernel_size self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormConv2d ( in_channels = channels [ i ], out_channels = channels [ i + 1 ], kernel_size = kernel_size [ i ], stride = 1 , padding = kernel_size [ i ] // 2 , bias = True , coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers ) def forward ( self , x ): return self . net ( x )","title":"LipschitzCNN"},{"location":"references/#normflows.nets.lipschitz.LipschitzCNN.__init__","text":"Constructor Parameters: Name Type Description Default channels Integer list with the number of channels of the layers required kernel_size Integer list of kernel sizes of the layers required lipschitz_const Maximum Lipschitz constant of each layer 0.97 max_lipschitz_iter Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used 5 lipschitz_tolerance Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 None init_zeros Flag, whether to initialize last layer approximately with zeros True Source code in normflows/nets/lipschitz.py 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 def __init__ ( self , channels , kernel_size , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\"Constructor Args: channels: Integer list with the number of channels of the layers kernel_size: Integer list of kernel sizes of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( kernel_size ) self . channels = channels self . kernel_size = kernel_size self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormConv2d ( in_channels = channels [ i ], out_channels = channels [ i + 1 ], kernel_size = kernel_size [ i ], stride = 1 , padding = kernel_size [ i ] // 2 , bias = True , coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers )","title":"__init__()"},{"location":"references/#normflows.nets.lipschitz.LipschitzMLP","text":"Bases: nn . Module Fully connected neural net which is Lipschitz continuou with Lipschitz constant L < 1 Source code in normflows/nets/lipschitz.py 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 class LipschitzMLP ( nn . Module ): \"\"\"Fully connected neural net which is Lipschitz continuou with Lipschitz constant L < 1\"\"\" def __init__ ( self , channels , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\" Constructor channels: Integer list with the number of channels of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( channels ) - 1 self . channels = channels self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormLinear ( in_features = channels [ i ], out_features = channels [ i + 1 ], coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers ) def forward ( self , x ): return self . net ( x )","title":"LipschitzMLP"},{"location":"references/#normflows.nets.lipschitz.LipschitzMLP.__init__","text":"Constructor channels: Integer list with the number of channels of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros Source code in normflows/nets/lipschitz.py 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 def __init__ ( self , channels , lipschitz_const = 0.97 , max_lipschitz_iter = 5 , lipschitz_tolerance = None , init_zeros = True , ): \"\"\" Constructor channels: Integer list with the number of channels of the layers lipschitz_const: Maximum Lipschitz constant of each layer max_lipschitz_iter: Maximum number of iterations used to ensure that layers are Lipschitz continuous with L smaller than set maximum; if None, tolerance is used lipschitz_tolerance: Float, tolerance used to ensure Lipschitz continuity if max_lipschitz_iter is None, typically 1e-3 init_zeros: Flag, whether to initialize last layer approximately with zeros \"\"\" super () . __init__ () self . n_layers = len ( channels ) - 1 self . channels = channels self . lipschitz_const = lipschitz_const self . max_lipschitz_iter = max_lipschitz_iter self . lipschitz_tolerance = lipschitz_tolerance self . init_zeros = init_zeros layers = [] for i in range ( self . n_layers ): layers += [ Swish (), InducedNormLinear ( in_features = channels [ i ], out_features = channels [ i + 1 ], coeff = lipschitz_const , domain = 2 , codomain = 2 , n_iterations = max_lipschitz_iter , atol = lipschitz_tolerance , rtol = lipschitz_tolerance , zero_init = init_zeros if i == ( self . n_layers - 1 ) else False , ), ] self . net = nn . Sequential ( * layers )","title":"__init__()"},{"location":"references/#normflows.nets.lipschitz.projmax_","text":"Inplace argmax on absolute value. Source code in normflows/nets/lipschitz.py 651 652 653 654 655 656 def projmax_ ( v ): \"\"\"Inplace argmax on absolute value.\"\"\" ind = torch . argmax ( torch . abs ( v )) v . zero_ () v [ ind ] = 1 return v","title":"projmax_()"},{"location":"references/#normflows.nets.made","text":"Implementation of MADE. Code taken from https://github.com/bayesiains/nsf","title":"made"},{"location":"references/#normflows.nets.made.MADE","text":"Bases: nn . Module Implementation of MADE. It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no). Source code in normflows/nets/made.py 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 class MADE ( nn . Module ): \"\"\"Implementation of MADE. It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no). \"\"\" def __init__ ( self , features , hidden_features , context_features = None , num_blocks = 2 , output_multiplier = 1 , use_residual_blocks = True , random_mask = False , permute_mask = False , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , preprocessing = None , ): if use_residual_blocks and random_mask : raise ValueError ( \"Residual blocks can't be used with random masks.\" ) super () . __init__ () # Preprocessing if preprocessing is None : self . preprocessing = lambda inputs : inputs else : self . preprocessing = preprocessing # Initial layer. input_degrees_ = _get_input_degrees ( features ) if permute_mask : input_degrees_ = input_degrees_ [ torch . randperm ( features )] self . initial_layer = MaskedLinear ( in_degrees = input_degrees_ , out_features = hidden_features , autoregressive_features = features , random_mask = random_mask , is_output = False , ) if context_features is not None : self . context_layer = nn . Linear ( context_features , hidden_features ) # Residual blocks. blocks = [] if use_residual_blocks : block_constructor = MaskedResidualBlock else : block_constructor = MaskedFeedforwardBlock prev_out_degrees = self . initial_layer . degrees for _ in range ( num_blocks ): blocks . append ( block_constructor ( in_degrees = prev_out_degrees , autoregressive_features = features , context_features = context_features , random_mask = random_mask , activation = activation , dropout_probability = dropout_probability , use_batch_norm = use_batch_norm , ) ) prev_out_degrees = blocks [ - 1 ] . degrees self . blocks = nn . ModuleList ( blocks ) # Final layer. self . final_layer = MaskedLinear ( in_degrees = prev_out_degrees , out_features = features * output_multiplier , autoregressive_features = features , random_mask = random_mask , is_output = True , out_degrees_ = input_degrees_ , ) def forward ( self , inputs , context = None ): outputs = self . preprocessing ( inputs ) outputs = self . initial_layer ( outputs ) if context is not None : outputs += self . context_layer ( context ) for block in self . blocks : outputs = block ( outputs , context ) outputs = self . final_layer ( outputs ) return outputs","title":"MADE"},{"location":"references/#normflows.nets.made.MaskedFeedforwardBlock","text":"Bases: nn . Module A feedforward block based on a masked linear module. NOTE In this implementation, the number of output features is taken to be equal to the number of input features. Source code in normflows/nets/made.py 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 class MaskedFeedforwardBlock ( nn . Module ): \"\"\"A feedforward block based on a masked linear module. **NOTE** In this implementation, the number of output features is taken to be equal to the number of input features. \"\"\" def __init__ ( self , in_degrees , autoregressive_features , context_features = None , random_mask = False , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , ): super () . __init__ () features = len ( in_degrees ) # Batch norm. if use_batch_norm : self . batch_norm = nn . BatchNorm1d ( features , eps = 1e-3 ) else : self . batch_norm = None if context_features is not None : raise NotImplementedError () # Masked linear. self . linear = MaskedLinear ( in_degrees = in_degrees , out_features = features , autoregressive_features = autoregressive_features , random_mask = random_mask , is_output = False , ) self . degrees = self . linear . degrees # Activation and dropout. self . activation = activation self . dropout = nn . Dropout ( p = dropout_probability ) def forward ( self , inputs , context = None ): if context is not None : raise NotImplementedError () if self . batch_norm : outputs = self . batch_norm ( inputs ) else : outputs = inputs outputs = self . linear ( outputs ) outputs = self . activation ( outputs ) outputs = self . dropout ( outputs ) return outputs","title":"MaskedFeedforwardBlock"},{"location":"references/#normflows.nets.made.MaskedLinear","text":"Bases: nn . Linear A linear module with a masked weight matrix. Source code in normflows/nets/made.py 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 class MaskedLinear ( nn . Linear ): \"\"\"A linear module with a masked weight matrix.\"\"\" def __init__ ( self , in_degrees , out_features , autoregressive_features , random_mask , is_output , bias = True , out_degrees_ = None , ): super () . __init__ ( in_features = len ( in_degrees ), out_features = out_features , bias = bias ) mask , degrees = self . _get_mask_and_degrees ( in_degrees = in_degrees , out_features = out_features , autoregressive_features = autoregressive_features , random_mask = random_mask , is_output = is_output , out_degrees_ = out_degrees_ , ) self . register_buffer ( \"mask\" , mask ) self . register_buffer ( \"degrees\" , degrees ) @classmethod def _get_mask_and_degrees ( cls , in_degrees , out_features , autoregressive_features , random_mask , is_output , out_degrees_ = None , ): if is_output : if out_degrees_ is None : out_degrees_ = _get_input_degrees ( autoregressive_features ) out_degrees = tile ( out_degrees_ , out_features // autoregressive_features ) mask = ( out_degrees [ ... , None ] > in_degrees ) . float () else : if random_mask : min_in_degree = torch . min ( in_degrees ) . item () min_in_degree = min ( min_in_degree , autoregressive_features - 1 ) out_degrees = torch . randint ( low = min_in_degree , high = autoregressive_features , size = [ out_features ], dtype = torch . long , ) else : max_ = max ( 1 , autoregressive_features - 1 ) min_ = min ( 1 , autoregressive_features - 1 ) out_degrees = torch . arange ( out_features ) % max_ + min_ mask = ( out_degrees [ ... , None ] >= in_degrees ) . float () return mask , out_degrees def forward ( self , x ): return F . linear ( x , self . weight * self . mask , self . bias )","title":"MaskedLinear"},{"location":"references/#normflows.nets.made.MaskedResidualBlock","text":"Bases: nn . Module A residual block containing masked linear modules. Source code in normflows/nets/made.py 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 class MaskedResidualBlock ( nn . Module ): \"\"\"A residual block containing masked linear modules.\"\"\" def __init__ ( self , in_degrees , autoregressive_features , context_features = None , random_mask = False , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , zero_initialization = True , ): if random_mask : raise ValueError ( \"Masked residual block can't be used with random masks.\" ) super () . __init__ () features = len ( in_degrees ) if context_features is not None : self . context_layer = nn . Linear ( context_features , features ) # Batch norm. self . use_batch_norm = use_batch_norm if use_batch_norm : self . batch_norm_layers = nn . ModuleList ( [ nn . BatchNorm1d ( features , eps = 1e-3 ) for _ in range ( 2 )] ) # Masked linear. linear_0 = MaskedLinear ( in_degrees = in_degrees , out_features = features , autoregressive_features = autoregressive_features , random_mask = False , is_output = False , ) linear_1 = MaskedLinear ( in_degrees = linear_0 . degrees , out_features = features , autoregressive_features = autoregressive_features , random_mask = False , is_output = False , ) self . linear_layers = nn . ModuleList ([ linear_0 , linear_1 ]) self . degrees = linear_1 . degrees if torch . all ( self . degrees >= in_degrees ) . item () != 1 : raise RuntimeError ( \"In a masked residual block, the output degrees can't be\" \" less than the corresponding input degrees.\" ) # Activation and dropout self . activation = activation self . dropout = nn . Dropout ( p = dropout_probability ) # Initialization. if zero_initialization : init . uniform_ ( self . linear_layers [ - 1 ] . weight , a =- 1e-3 , b = 1e-3 ) init . uniform_ ( self . linear_layers [ - 1 ] . bias , a =- 1e-3 , b = 1e-3 ) def forward ( self , inputs , context = None ): temps = inputs if self . use_batch_norm : temps = self . batch_norm_layers [ 0 ]( temps ) temps = self . activation ( temps ) temps = self . linear_layers [ 0 ]( temps ) if self . use_batch_norm : temps = self . batch_norm_layers [ 1 ]( temps ) temps = self . activation ( temps ) temps = self . dropout ( temps ) temps = self . linear_layers [ 1 ]( temps ) if context is not None : temps = F . glu ( torch . cat (( temps , self . context_layer ( context )), dim = 1 ), dim = 1 ) return inputs + temps","title":"MaskedResidualBlock"},{"location":"references/#normflows.nets.made_test","text":"Tests for MADE. Code partially taken from https://github.com/bayesiains/nsf","title":"made_test"},{"location":"references/#normflows.nets.mlp","text":"","title":"mlp"},{"location":"references/#normflows.nets.mlp.MLP","text":"Bases: nn . Module A multilayer perceptron with Leaky ReLU nonlinearities Source code in normflows/nets/mlp.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 class MLP ( nn . Module ): \"\"\" A multilayer perceptron with Leaky ReLU nonlinearities \"\"\" def __init__ ( self , layers , leaky = 0.0 , score_scale = None , output_fn = None , output_scale = None , init_zeros = False , dropout = None , ): \"\"\" layers: list of layer sizes from start to end leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used score_scale: Factor to apply to the scores, i.e. output before output_fn. output_fn: String, function to be applied to the output, either None, \"sigmoid\", \"relu\", \"tanh\", or \"clampexp\" output_scale: Rescale outputs if output_fn is specified, i.e. ```scale * output_fn(out / scale)``` init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see [arXiv 1807.03039](https://arxiv.org/abs/1807.03039)) dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done \"\"\" super () . __init__ () net = nn . ModuleList ([]) for k in range ( len ( layers ) - 2 ): net . append ( nn . Linear ( layers [ k ], layers [ k + 1 ])) net . append ( nn . LeakyReLU ( leaky )) if dropout is not None : net . append ( nn . Dropout ( p = dropout )) net . append ( nn . Linear ( layers [ - 2 ], layers [ - 1 ])) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) if output_fn is not None : if score_scale is not None : net . append ( utils . ConstScaleLayer ( score_scale )) if output_fn == \"sigmoid\" : net . append ( nn . Sigmoid ()) elif output_fn == \"relu\" : net . append ( nn . ReLU ()) elif output_fn == \"tanh\" : net . append ( nn . Tanh ()) elif output_fn == \"clampexp\" : net . append ( utils . ClampExp ()) else : NotImplementedError ( \"This output function is not implemented.\" ) if output_scale is not None : net . append ( utils . ConstScaleLayer ( output_scale )) self . net = nn . Sequential ( * net ) def forward ( self , x ): return self . net ( x )","title":"MLP"},{"location":"references/#normflows.nets.mlp.MLP.__init__","text":"layers: list of layer sizes from start to end leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used score_scale: Factor to apply to the scores, i.e. output before output_fn. output_fn: String, function to be applied to the output, either None, \"sigmoid\", \"relu\", \"tanh\", or \"clampexp\" output_scale: Rescale outputs if output_fn is specified, i.e. scale * output_fn(out / scale) init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see arXiv 1807.03039 ) dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done Source code in normflows/nets/mlp.py 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def __init__ ( self , layers , leaky = 0.0 , score_scale = None , output_fn = None , output_scale = None , init_zeros = False , dropout = None , ): \"\"\" layers: list of layer sizes from start to end leaky: slope of the leaky part of the ReLU, if 0.0, standard ReLU is used score_scale: Factor to apply to the scores, i.e. output before output_fn. output_fn: String, function to be applied to the output, either None, \"sigmoid\", \"relu\", \"tanh\", or \"clampexp\" output_scale: Rescale outputs if output_fn is specified, i.e. ```scale * output_fn(out / scale)``` init_zeros: Flag, if true, weights and biases of last layer are initialized with zeros (helpful for deep models, see [arXiv 1807.03039](https://arxiv.org/abs/1807.03039)) dropout: Float, if specified, dropout is done before last layer; if None, no dropout is done \"\"\" super () . __init__ () net = nn . ModuleList ([]) for k in range ( len ( layers ) - 2 ): net . append ( nn . Linear ( layers [ k ], layers [ k + 1 ])) net . append ( nn . LeakyReLU ( leaky )) if dropout is not None : net . append ( nn . Dropout ( p = dropout )) net . append ( nn . Linear ( layers [ - 2 ], layers [ - 1 ])) if init_zeros : nn . init . zeros_ ( net [ - 1 ] . weight ) nn . init . zeros_ ( net [ - 1 ] . bias ) if output_fn is not None : if score_scale is not None : net . append ( utils . ConstScaleLayer ( score_scale )) if output_fn == \"sigmoid\" : net . append ( nn . Sigmoid ()) elif output_fn == \"relu\" : net . append ( nn . ReLU ()) elif output_fn == \"tanh\" : net . append ( nn . Tanh ()) elif output_fn == \"clampexp\" : net . append ( utils . ClampExp ()) else : NotImplementedError ( \"This output function is not implemented.\" ) if output_scale is not None : net . append ( utils . ConstScaleLayer ( output_scale )) self . net = nn . Sequential ( * net )","title":"__init__()"},{"location":"references/#normflows.nets.resnet","text":"","title":"resnet"},{"location":"references/#normflows.nets.resnet.ResidualBlock","text":"Bases: nn . Module A general-purpose residual block. Works only with 1-dim inputs. Source code in normflows/nets/resnet.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 class ResidualBlock ( nn . Module ): \"\"\"A general-purpose residual block. Works only with 1-dim inputs.\"\"\" def __init__ ( self , features , context_features , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , zero_initialization = True , ): super () . __init__ () self . activation = activation self . use_batch_norm = use_batch_norm if use_batch_norm : self . batch_norm_layers = nn . ModuleList ( [ nn . BatchNorm1d ( features , eps = 1e-3 ) for _ in range ( 2 )] ) if context_features is not None : self . context_layer = nn . Linear ( context_features , features ) self . linear_layers = nn . ModuleList ( [ nn . Linear ( features , features ) for _ in range ( 2 )] ) self . dropout = nn . Dropout ( p = dropout_probability ) if zero_initialization : init . uniform_ ( self . linear_layers [ - 1 ] . weight , - 1e-3 , 1e-3 ) init . uniform_ ( self . linear_layers [ - 1 ] . bias , - 1e-3 , 1e-3 ) def forward ( self , inputs , context = None ): temps = inputs if self . use_batch_norm : temps = self . batch_norm_layers [ 0 ]( temps ) temps = self . activation ( temps ) temps = self . linear_layers [ 0 ]( temps ) if self . use_batch_norm : temps = self . batch_norm_layers [ 1 ]( temps ) temps = self . activation ( temps ) temps = self . dropout ( temps ) temps = self . linear_layers [ 1 ]( temps ) if context is not None : temps = F . glu ( torch . cat (( temps , self . context_layer ( context )), dim = 1 ), dim = 1 ) return inputs + temps","title":"ResidualBlock"},{"location":"references/#normflows.nets.resnet.ResidualNet","text":"Bases: nn . Module A general-purpose residual network. Works only with 1-dim inputs. Source code in normflows/nets/resnet.py 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 class ResidualNet ( nn . Module ): \"\"\"A general-purpose residual network. Works only with 1-dim inputs.\"\"\" def __init__ ( self , in_features , out_features , hidden_features , context_features = None , num_blocks = 2 , activation = F . relu , dropout_probability = 0.0 , use_batch_norm = False , preprocessing = None , ): super () . __init__ () self . hidden_features = hidden_features self . context_features = context_features self . preprocessing = preprocessing if context_features is not None : self . initial_layer = nn . Linear ( in_features + context_features , hidden_features ) else : self . initial_layer = nn . Linear ( in_features , hidden_features ) self . blocks = nn . ModuleList ( [ ResidualBlock ( features = hidden_features , context_features = context_features , activation = activation , dropout_probability = dropout_probability , use_batch_norm = use_batch_norm , ) for _ in range ( num_blocks ) ] ) self . final_layer = nn . Linear ( hidden_features , out_features ) def forward ( self , inputs , context = None ): if self . preprocessing is None : temps = inputs else : temps = self . preprocessing ( inputs ) if context is None : temps = self . initial_layer ( temps ) else : temps = self . initial_layer ( torch . cat (( temps , context ), dim = 1 )) for block in self . blocks : temps = block ( temps , context = context ) outputs = self . final_layer ( temps ) return outputs","title":"ResidualNet"},{"location":"references/#normflows.sampling","text":"","title":"sampling"},{"location":"references/#normflows.sampling.hais","text":"","title":"hais"},{"location":"references/#normflows.sampling.hais.HAIS","text":"Class which performs HAIS Source code in normflows/sampling/hais.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 class HAIS : \"\"\" Class which performs HAIS \"\"\" def __init__ ( self , betas , prior , target , num_leapfrog , step_size , log_mass ): \"\"\" Args: betas: Annealing schedule, the jth target is ```f_j(x) = f_0(x)^{\\beta_j} f_n(x)^{1-\\beta_j}``` where the target is proportional to f_0 and the prior is proportional to f_n. The number of intermediate steps is infered from the shape of betas. Should be of the form 1 = \\beta_0 > \\beta_1 > ... > \\beta_n = 0 prior: The prior distribution to start the HAIS chain. target: The target distribution from which we would like to draw weighted samples. num_leapfrog: Number of leapfrog steps in the HMC transitions. step_size: step_size to use for HMC transitions. log_mass: log_mass to use for HMC transitions. \"\"\" self . prior = prior self . target = target self . layers = [] n = betas . shape [ 0 ] - 1 for i in range ( n - 1 , 0 , - 1 ): intermediate_target = distributions . LinearInterpolation ( self . target , self . prior , betas [ i ] ) self . layers += [ flows . HamiltonianMonteCarlo ( intermediate_target , num_leapfrog , torch . log ( step_size ), log_mass ) ] def sample ( self , num_samples ): \"\"\"Run HAIS to draw samples from the target with appropriate weights. Args: num_samples: The number of samples to draw.a \"\"\" samples , log_weights = self . prior . forward ( num_samples ) log_weights = - log_weights for i in range ( len ( self . layers )): samples , log_weights_addition = self . layers [ i ] . forward ( samples ) log_weights += log_weights_addition log_weights += self . target . log_prob ( samples ) return samples , log_weights","title":"HAIS"},{"location":"references/#normflows.sampling.hais.HAIS.__init__","text":"Parameters: Name Type Description Default betas Annealing schedule, the jth target is f_j(x) = f_0(x)^{\beta_j} f_n(x)^{1-\beta_j} where the target is proportional to f_0 and the prior is proportional to f_n. The number of intermediate steps is infered from the shape of betas. Should be of the form 1 = \beta_0 > \beta_1 > ... > \beta_n = 0 required prior The prior distribution to start the HAIS chain. required target The target distribution from which we would like to draw weighted samples. required num_leapfrog Number of leapfrog steps in the HMC transitions. required step_size step_size to use for HMC transitions. required log_mass log_mass to use for HMC transitions. required Source code in normflows/sampling/hais.py 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 def __init__ ( self , betas , prior , target , num_leapfrog , step_size , log_mass ): \"\"\" Args: betas: Annealing schedule, the jth target is ```f_j(x) = f_0(x)^{\\beta_j} f_n(x)^{1-\\beta_j}``` where the target is proportional to f_0 and the prior is proportional to f_n. The number of intermediate steps is infered from the shape of betas. Should be of the form 1 = \\beta_0 > \\beta_1 > ... > \\beta_n = 0 prior: The prior distribution to start the HAIS chain. target: The target distribution from which we would like to draw weighted samples. num_leapfrog: Number of leapfrog steps in the HMC transitions. step_size: step_size to use for HMC transitions. log_mass: log_mass to use for HMC transitions. \"\"\" self . prior = prior self . target = target self . layers = [] n = betas . shape [ 0 ] - 1 for i in range ( n - 1 , 0 , - 1 ): intermediate_target = distributions . LinearInterpolation ( self . target , self . prior , betas [ i ] ) self . layers += [ flows . HamiltonianMonteCarlo ( intermediate_target , num_leapfrog , torch . log ( step_size ), log_mass ) ]","title":"__init__()"},{"location":"references/#normflows.sampling.hais.HAIS.sample","text":"Run HAIS to draw samples from the target with appropriate weights. Parameters: Name Type Description Default num_samples The number of samples to draw.a required Source code in normflows/sampling/hais.py 37 38 39 40 41 42 43 44 45 46 47 48 49 def sample ( self , num_samples ): \"\"\"Run HAIS to draw samples from the target with appropriate weights. Args: num_samples: The number of samples to draw.a \"\"\" samples , log_weights = self . prior . forward ( num_samples ) log_weights = - log_weights for i in range ( len ( self . layers )): samples , log_weights_addition = self . layers [ i ] . forward ( samples ) log_weights += log_weights_addition log_weights += self . target . log_prob ( samples ) return samples , log_weights","title":"sample()"},{"location":"references/#normflows.transforms","text":"","title":"transforms"},{"location":"references/#normflows.transforms.Logit","text":"Bases: flows . Flow Logit mapping of image tensor, see RealNVP paper logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) Source code in normflows/transforms.py 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 class Logit ( flows . Flow ): \"\"\"Logit mapping of image tensor, see RealNVP paper ``` logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) ``` \"\"\" def __init__ ( self , alpha = 0.05 ): \"\"\"Constructor Args: alpha: Alpha parameter, see above \"\"\" super () . __init__ () self . alpha = alpha def forward ( self , z ): beta = 1 - 2 * self . alpha sum_dims = list ( range ( 1 , z . dim ())) ls = torch . sum ( torch . nn . functional . logsigmoid ( z ), dim = sum_dims ) mls = torch . sum ( torch . nn . functional . logsigmoid ( - z ), dim = sum_dims ) log_det = - np . log ( beta ) * np . prod ([ * z . shape [ 1 :]]) + ls + mls z = ( torch . sigmoid ( z ) - self . alpha ) / beta return z , log_det def inverse ( self , z ): beta = 1 - 2 * self . alpha z = self . alpha + beta * z logz = torch . log ( z ) log1mz = torch . log ( 1 - z ) z = logz - log1mz sum_dims = list ( range ( 1 , z . dim ())) log_det = ( np . log ( beta ) * np . prod ([ * z . shape [ 1 :]]) - torch . sum ( logz , dim = sum_dims ) - torch . sum ( log1mz , dim = sum_dims ) ) return z , log_det","title":"Logit"},{"location":"references/#normflows.transforms.Logit.__init__","text":"Constructor Parameters: Name Type Description Default alpha Alpha parameter, see above 0.05 Source code in normflows/transforms.py 17 18 19 20 21 22 23 24 def __init__ ( self , alpha = 0.05 ): \"\"\"Constructor Args: alpha: Alpha parameter, see above \"\"\" super () . __init__ () self . alpha = alpha","title":"__init__()"},{"location":"references/#normflows.transforms.Shift","text":"Bases: flows . Flow Shift data by a fixed constant Default is -0.5 to shift data from interval [0, 1] to [-0.5, 0.5] Source code in normflows/transforms.py 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 class Shift ( flows . Flow ): \"\"\"Shift data by a fixed constant Default is -0.5 to shift data from interval [0, 1] to [-0.5, 0.5] \"\"\" def __init__ ( self , shift =- 0.5 ): \"\"\"Constructor Args: shift: Shift to apply to the data \"\"\" super () . __init__ () self . shift = shift def forward ( self , z ): z -= self . shift log_det = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) return z , log_det def inverse ( self , z ): z += self . shift log_det = torch . zeros ( z . shape [ 0 ], dtype = z . dtype , device = z . device ) return z , log_det","title":"Shift"},{"location":"references/#normflows.transforms.Shift.__init__","text":"Constructor Parameters: Name Type Description Default shift Shift to apply to the data -0.5 Source code in normflows/transforms.py 57 58 59 60 61 62 63 64 def __init__ ( self , shift =- 0.5 ): \"\"\"Constructor Args: shift: Shift to apply to the data \"\"\" super () . __init__ () self . shift = shift","title":"__init__()"},{"location":"references/#normflows.transforms_test","text":"","title":"transforms_test"},{"location":"references/#normflows.utils","text":"","title":"utils"},{"location":"references/#normflows.utils.eval","text":"","title":"eval"},{"location":"references/#normflows.utils.eval.bitsPerDim","text":"Computes the bits per dim for a batch of data Parameters: Name Type Description Default model Model to compute bits per dim for required x Batch of data required y Class labels for batch of data if base distribution is class conditional None trans Transformation to be applied to images during training 'logit' trans_param List of parameters of the transformation [0.05] Returns: Type Description Bits per dim for data batch under model Source code in normflows/utils/eval.py 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 def bitsPerDim ( model , x , y = None , trans = \"logit\" , trans_param = [ 0.05 ]): \"\"\"Computes the bits per dim for a batch of data Args: model: Model to compute bits per dim for x: Batch of data y: Class labels for batch of data if base distribution is class conditional trans: Transformation to be applied to images during training trans_param: List of parameters of the transformation Returns: Bits per dim for data batch under model \"\"\" dims = torch . prod ( torch . tensor ( x . size ()[ 1 :])) if trans == \"logit\" : if y is None : log_q = model . log_prob ( x ) else : log_q = model . log_prob ( x , y ) sum_dims = list ( range ( 1 , x . dim ())) ls = torch . nn . LogSigmoid () sig_ = torch . sum ( ls ( x ) / np . log ( 2 ), sum_dims ) sig_ += torch . sum ( ls ( - x ) / np . log ( 2 ), sum_dims ) b = - log_q / dims / np . log ( 2 ) - np . log2 ( 1 - trans_param [ 0 ]) + 8 b += sig_ / dims else : raise NotImplementedError ( \"The transformation \" + trans + \" is not implemented.\" ) return b","title":"bitsPerDim()"},{"location":"references/#normflows.utils.eval.bitsPerDimDataset","text":"Computes average bits per dim for an entire dataset given by a data loader Parameters: Name Type Description Default model Model to compute bits per dim for required data_loader Data loader of dataset required class_cond Flag indicating whether model is class_conditional True trans Transformation to be applied to images during training 'logit' trans_param List of parameters of the transformation [0.05] Returns: Type Description Average bits per dim for dataset Source code in normflows/utils/eval.py 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 def bitsPerDimDataset ( model , data_loader , class_cond = True , trans = \"logit\" , trans_param = [ 0.05 ] ): \"\"\"Computes average bits per dim for an entire dataset given by a data loader Args: model: Model to compute bits per dim for data_loader: Data loader of dataset class_cond: Flag indicating whether model is class_conditional trans: Transformation to be applied to images during training trans_param: List of parameters of the transformation Returns: Average bits per dim for dataset \"\"\" n = 0 b_cum = 0 with torch . no_grad (): for x , y in iter ( data_loader ): b_ = bitsPerDim ( model , x , y . to ( x . device ) if class_cond else None , trans , trans_param ) b_np = b_ . to ( \"cpu\" ) . numpy () b_cum += np . nansum ( b_np ) n += len ( x ) - np . sum ( np . isnan ( b_np )) b = b_cum / n return b","title":"bitsPerDimDataset()"},{"location":"references/#normflows.utils.masks","text":"","title":"masks"},{"location":"references/#normflows.utils.masks.create_alternating_binary_mask","text":"Creates a binary mask of a given dimension which alternates its masking. Parameters: Name Type Description Default features Dimension of mask. required even If True, even values are assigned 1s, odd 0s. If False, vice versa. True Returns: Type Description Alternating binary mask of type torch.Tensor. Source code in normflows/utils/masks.py 4 5 6 7 8 9 10 11 12 13 14 15 16 17 def create_alternating_binary_mask ( features , even = True ): \"\"\"Creates a binary mask of a given dimension which alternates its masking. Args: features: Dimension of mask. even: If True, even values are assigned 1s, odd 0s. If False, vice versa. Returns: Alternating binary mask of type torch.Tensor. \"\"\" mask = torch . zeros ( features ) . byte () start = 0 if even else 1 mask [ start :: 2 ] += 1 return mask","title":"create_alternating_binary_mask()"},{"location":"references/#normflows.utils.masks.create_mid_split_binary_mask","text":"Creates a binary mask of a given dimension which splits its masking at the midpoint. Parameters: Name Type Description Default features Dimension of mask. required Returns: Type Description Binary mask split at midpoint of type torch.Tensor Source code in normflows/utils/masks.py 20 21 22 23 24 25 26 27 28 29 30 31 32 def create_mid_split_binary_mask ( features ): \"\"\"Creates a binary mask of a given dimension which splits its masking at the midpoint. Args: features: Dimension of mask. Returns: Binary mask split at midpoint of type torch.Tensor \"\"\" mask = torch . zeros ( features ) . byte () midpoint = features // 2 if features % 2 == 0 else features // 2 + 1 mask [: midpoint ] += 1 return mask","title":"create_mid_split_binary_mask()"},{"location":"references/#normflows.utils.masks.create_random_binary_mask","text":"Creates a random binary mask of a given dimension with half of its entries randomly set to 1s. Parameters: Name Type Description Default features Dimension of mask. required seed Seed to be used None Returns: Type Description Binary mask with half of its entries set to 1s, of type torch.Tensor. Source code in normflows/utils/masks.py 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 def create_random_binary_mask ( features , seed = None ): \"\"\"Creates a random binary mask of a given dimension with half of its entries randomly set to 1s. Args: features: Dimension of mask. seed: Seed to be used Returns: Binary mask with half of its entries set to 1s, of type torch.Tensor. \"\"\" mask = torch . zeros ( features ) . byte () weights = torch . ones ( features ) . float () num_samples = features // 2 if features % 2 == 0 else features // 2 + 1 if seed is None : generator = None else : generator = torch . Generator () generator . manual_seed ( seed ) indices = torch . multinomial ( input = weights , num_samples = num_samples , replacement = False , generator = generator ) mask [ indices ] += 1 return mask","title":"create_random_binary_mask()"},{"location":"references/#normflows.utils.nn","text":"","title":"nn"},{"location":"references/#normflows.utils.nn.ActNorm","text":"Bases: nn . Module ActNorm layer with just one forward pass Source code in normflows/utils/nn.py 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 class ActNorm ( nn . Module ): \"\"\" ActNorm layer with just one forward pass \"\"\" def __init__ ( self , shape ): \"\"\"Constructor Args: shape: Same as shape in flows.ActNorm logscale_factor: Same as shape in flows.ActNorm \"\"\" super () . __init__ () self . actNorm = flows . ActNorm ( shape ) def forward ( self , input ): out , _ = self . actNorm ( input ) return out","title":"ActNorm"},{"location":"references/#normflows.utils.nn.ActNorm.__init__","text":"Constructor Parameters: Name Type Description Default shape Same as shape in flows.ActNorm required logscale_factor Same as shape in flows.ActNorm required Source code in normflows/utils/nn.py 30 31 32 33 34 35 36 37 38 39 def __init__ ( self , shape ): \"\"\"Constructor Args: shape: Same as shape in flows.ActNorm logscale_factor: Same as shape in flows.ActNorm \"\"\" super () . __init__ () self . actNorm = flows . ActNorm ( shape )","title":"__init__()"},{"location":"references/#normflows.utils.nn.ClampExp","text":"Bases: nn . Module Nonlinearity min(exp(lam * x), 1) Source code in normflows/utils/nn.py 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 class ClampExp ( nn . Module ): \"\"\" Nonlinearity min(exp(lam * x), 1) \"\"\" def __init__ ( self ): \"\"\"Constructor Args: lam: Lambda parameter \"\"\" super ( ClampExp , self ) . __init__ () def forward ( self , x ): one = torch . tensor ( 1.0 , device = x . device , dtype = x . dtype ) return torch . min ( torch . exp ( x ), one )","title":"ClampExp"},{"location":"references/#normflows.utils.nn.ClampExp.__init__","text":"Constructor Parameters: Name Type Description Default lam Lambda parameter required Source code in normflows/utils/nn.py 51 52 53 54 55 56 57 def __init__ ( self ): \"\"\"Constructor Args: lam: Lambda parameter \"\"\" super ( ClampExp , self ) . __init__ ()","title":"__init__()"},{"location":"references/#normflows.utils.nn.ConstScaleLayer","text":"Bases: nn . Module Scaling features by a fixed factor Source code in normflows/utils/nn.py 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 class ConstScaleLayer ( nn . Module ): \"\"\" Scaling features by a fixed factor \"\"\" def __init__ ( self , scale = 1.0 ): \"\"\"Constructor Args: scale: Scale to apply to features \"\"\" super () . __init__ () self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu ) def forward ( self , input ): return input * self . scale","title":"ConstScaleLayer"},{"location":"references/#normflows.utils.nn.ConstScaleLayer.__init__","text":"Constructor Parameters: Name Type Description Default scale Scale to apply to features 1.0 Source code in normflows/utils/nn.py 12 13 14 15 16 17 18 19 20 def __init__ ( self , scale = 1.0 ): \"\"\"Constructor Args: scale: Scale to apply to features \"\"\" super () . __init__ () self . scale_cpu = torch . tensor ( scale ) self . register_buffer ( \"scale\" , self . scale_cpu )","title":"__init__()"},{"location":"references/#normflows.utils.nn.PeriodicFeaturesCat","text":"Bases: nn . Module Converts a specified part of the input to periodic features by replacing those features f with [sin(scale * f), cos(scale * f)]. Note that this decreases the number of features and their order is changed. Source code in normflows/utils/nn.py 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 class PeriodicFeaturesCat ( nn . Module ): \"\"\" Converts a specified part of the input to periodic features by replacing those features f with [sin(scale * f), cos(scale * f)]. Note that this decreases the number of features and their order is changed. \"\"\" def __init__ ( self , ndim , ind , scale = 1.0 ): \"\"\" Constructor :param ndim: Int, number of dimensions :param ind: Iterable, indices of input elements to convert to periodic features :param scale: Scalar or iterable, used to scale inputs before converting them to periodic features \"\"\" super ( PeriodicFeaturesCat , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale def forward ( self , inputs ): inputs_ = inputs [ ... , self . ind ] inputs_ = self . scale * inputs_ inputs_sin = torch . sin ( inputs_ ) inputs_cos = torch . cos ( inputs_ ) out = torch . cat (( inputs_sin , inputs_cos , inputs [ ... , self . ind_ ]), - 1 ) return out","title":"PeriodicFeaturesCat"},{"location":"references/#normflows.utils.nn.PeriodicFeaturesCat.__init__","text":"Constructor :param ndim: Int, number of dimensions :param ind: Iterable, indices of input elements to convert to periodic features :param scale: Scalar or iterable, used to scale inputs before converting them to periodic features Source code in normflows/utils/nn.py 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 def __init__ ( self , ndim , ind , scale = 1.0 ): \"\"\" Constructor :param ndim: Int, number of dimensions :param ind: Iterable, indices of input elements to convert to periodic features :param scale: Scalar or iterable, used to scale inputs before converting them to periodic features \"\"\" super ( PeriodicFeaturesCat , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale","title":"__init__()"},{"location":"references/#normflows.utils.nn.PeriodicFeaturesElementwise","text":"Bases: nn . Module Converts a specified part of the input to periodic features by replacing those features f with w1 * sin(scale * f) + w2 * cos(scale * f). Note that this operation is done elementwise and, therefore, some information about the feature can be lost. Source code in normflows/utils/nn.py 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 class PeriodicFeaturesElementwise ( nn . Module ): \"\"\" Converts a specified part of the input to periodic features by replacing those features f with w1 * sin(scale * f) + w2 * cos(scale * f). Note that this operation is done elementwise and, therefore, some information about the feature can be lost. \"\"\" def __init__ ( self , ndim , ind , scale = 1.0 , bias = False , activation = None ): \"\"\"Constructor Args: ndim (int): number of dimensions ind (iterable): indices of input elements to convert to periodic features scale: Scalar or iterable, used to scale inputs before converting them to periodic features bias: Flag, whether to add a bias activation: Function or None, activation function to be applied \"\"\" super ( PeriodicFeaturesElementwise , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) self . weights = nn . Parameter ( torch . ones ( len ( self . ind ), 2 )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale self . apply_bias = bias if self . apply_bias : self . bias = nn . Parameter ( torch . zeros ( len ( self . ind ))) if activation is None : self . activation = lambda input : input else : self . activation = activation def forward ( self , inputs ): inputs_ = inputs [ ... , self . ind ] inputs_ = self . scale * inputs_ inputs_ = self . weights [:, 0 ] * torch . sin ( inputs_ ) + self . weights [ :, 1 ] * torch . cos ( inputs_ ) if self . apply_bias : inputs_ = inputs_ + self . bias inputs_ = self . activation ( inputs_ ) out = torch . cat (( inputs_ , inputs [ ... , self . ind_ ]), - 1 ) return out [ ... , self . inv_perm ]","title":"PeriodicFeaturesElementwise"},{"location":"references/#normflows.utils.nn.PeriodicFeaturesElementwise.__init__","text":"Constructor Parameters: Name Type Description Default ndim int number of dimensions required ind iterable indices of input elements to convert to periodic features required scale Scalar or iterable, used to scale inputs before converting them to periodic features 1.0 bias Flag, whether to add a bias False activation Function or None, activation function to be applied None Source code in normflows/utils/nn.py 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 def __init__ ( self , ndim , ind , scale = 1.0 , bias = False , activation = None ): \"\"\"Constructor Args: ndim (int): number of dimensions ind (iterable): indices of input elements to convert to periodic features scale: Scalar or iterable, used to scale inputs before converting them to periodic features bias: Flag, whether to add a bias activation: Function or None, activation function to be applied \"\"\" super ( PeriodicFeaturesElementwise , self ) . __init__ () # Set up indices and permutations self . ndim = ndim if torch . is_tensor ( ind ): self . register_buffer ( \"ind\" , torch . _cast_Long ( ind )) else : self . register_buffer ( \"ind\" , torch . tensor ( ind , dtype = torch . long )) ind_ = [] for i in range ( self . ndim ): if not i in self . ind : ind_ += [ i ] self . register_buffer ( \"ind_\" , torch . tensor ( ind_ , dtype = torch . long )) perm_ = torch . cat (( self . ind , self . ind_ )) inv_perm_ = torch . zeros_like ( perm_ ) for i in range ( self . ndim ): inv_perm_ [ perm_ [ i ]] = i self . register_buffer ( \"inv_perm\" , inv_perm_ ) self . weights = nn . Parameter ( torch . ones ( len ( self . ind ), 2 )) if torch . is_tensor ( scale ): self . register_buffer ( \"scale\" , scale ) else : self . scale = scale self . apply_bias = bias if self . apply_bias : self . bias = nn . Parameter ( torch . zeros ( len ( self . ind ))) if activation is None : self . activation = lambda input : input else : self . activation = activation","title":"__init__()"},{"location":"references/#normflows.utils.nn.sum_except_batch","text":"Sums all elements of x except for the first num_batch_dims dimensions. Source code in normflows/utils/nn.py 190 191 192 193 def sum_except_batch ( x , num_batch_dims = 1 ): \"\"\"Sums all elements of `x` except for the first `num_batch_dims` dimensions.\"\"\" reduce_dims = list ( range ( num_batch_dims , x . ndimension ())) return torch . sum ( x , dim = reduce_dims )","title":"sum_except_batch()"},{"location":"references/#normflows.utils.optim","text":"","title":"optim"},{"location":"references/#normflows.utils.optim.clear_grad","text":"Set gradients of model parameter to None as this speeds up training, See youtube Parameters: Name Type Description Default model Model to clear gradients of required Source code in normflows/utils/optim.py 16 17 18 19 20 21 22 23 24 25 def clear_grad ( model ): \"\"\"Set gradients of model parameter to None as this speeds up training, See [youtube](https://www.youtube.com/watch?v=9mS1fIYj1So) Args: model: Model to clear gradients of \"\"\" for param in model . parameters (): param . grad = None","title":"clear_grad()"},{"location":"references/#normflows.utils.optim.set_requires_grad","text":"Sets requires_grad flag of all parameters of a torch.nn.module Parameters: Name Type Description Default module torch.nn.module required flag Flag to set requires_grad to required Source code in normflows/utils/optim.py 4 5 6 7 8 9 10 11 12 13 def set_requires_grad ( module , flag ): \"\"\"Sets requires_grad flag of all parameters of a torch.nn.module Args: module: torch.nn.module flag: Flag to set requires_grad to \"\"\" for param in module . parameters (): param . requires_grad = flag","title":"set_requires_grad()"},{"location":"references/#normflows.utils.preprocessing","text":"","title":"preprocessing"},{"location":"references/#normflows.utils.preprocessing.Jitter","text":"Transform for dataloader, adds uniform jitter noise to data Source code in normflows/utils/preprocessing.py 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 class Jitter : \"\"\"Transform for dataloader, adds uniform jitter noise to data\"\"\" def __init__ ( self , scale = 1.0 / 256 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale def __call__ ( self , x ): eps = torch . rand_like ( x ) * self . scale x_ = x + eps return x_","title":"Jitter"},{"location":"references/#normflows.utils.preprocessing.Jitter.__init__","text":"Constructor Parameters: Name Type Description Default scale Scaling factor for noise 1.0 / 256 Source code in normflows/utils/preprocessing.py 31 32 33 34 35 36 37 def __init__ ( self , scale = 1.0 / 256 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale","title":"__init__()"},{"location":"references/#normflows.utils.preprocessing.Logit","text":"Transform for dataloader logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) Source code in normflows/utils/preprocessing.py 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 class Logit : \"\"\"Transform for dataloader ``` logit(alpha + (1 - alpha) * x) where logit(x) = log(x / (1 - x)) ``` \"\"\" def __init__ ( self , alpha = 0 ): \"\"\"Constructor Args: alpha: see above \"\"\" self . alpha = alpha def __call__ ( self , x ): x_ = self . alpha + ( 1 - self . alpha ) * x return torch . log ( x_ / ( 1 - x_ )) def inverse ( self , x ): return ( torch . sigmoid ( x ) - self . alpha ) / ( 1 - self . alpha )","title":"Logit"},{"location":"references/#normflows.utils.preprocessing.Logit.__init__","text":"Constructor Parameters: Name Type Description Default alpha see above 0 Source code in normflows/utils/preprocessing.py 12 13 14 15 16 17 18 def __init__ ( self , alpha = 0 ): \"\"\"Constructor Args: alpha: see above \"\"\" self . alpha = alpha","title":"__init__()"},{"location":"references/#normflows.utils.preprocessing.Scale","text":"Transform for dataloader, adds uniform jitter noise to data Source code in normflows/utils/preprocessing.py 45 46 47 48 49 50 51 52 53 54 55 56 57 class Scale : \"\"\"Transform for dataloader, adds uniform jitter noise to data\"\"\" def __init__ ( self , scale = 255.0 / 256.0 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale def __call__ ( self , x ): return x * self . scale","title":"Scale"},{"location":"references/#normflows.utils.preprocessing.Scale.__init__","text":"Constructor Parameters: Name Type Description Default scale Scaling factor for noise 255.0 / 256.0 Source code in normflows/utils/preprocessing.py 48 49 50 51 52 53 54 def __init__ ( self , scale = 255.0 / 256.0 ): \"\"\"Constructor Args: scale: Scaling factor for noise \"\"\" self . scale = scale","title":"__init__()"}]} \ No newline at end of file diff --git a/search/worker.js b/search/worker.js new file mode 100644 index 0000000..8628dbc --- /dev/null +++ b/search/worker.js @@ -0,0 +1,133 @@ +var base_path = 'function' === typeof importScripts ? '.' : '/search/'; +var allowSearch = false; +var index; +var documents = {}; +var lang = ['en']; +var data; + +function getScript(script, callback) { + console.log('Loading script: ' + script); + $.getScript(base_path + script).done(function () { + callback(); + }).fail(function (jqxhr, settings, exception) { + console.log('Error: ' + exception); + }); +} + +function getScriptsInOrder(scripts, callback) { + if (scripts.length === 0) { + callback(); + return; + } + getScript(scripts[0], function() { + getScriptsInOrder(scripts.slice(1), callback); + }); +} + +function loadScripts(urls, callback) { + if( 'function' === typeof importScripts ) { + importScripts.apply(null, urls); + callback(); + } else { + getScriptsInOrder(urls, callback); + } +} + +function onJSONLoaded () { + data = JSON.parse(this.responseText); + var scriptsToLoad = ['lunr.js']; + if (data.config && data.config.lang && data.config.lang.length) { + lang = data.config.lang; + } + if (lang.length > 1 || lang[0] !== "en") { + scriptsToLoad.push('lunr.stemmer.support.js'); + if (lang.length > 1) { + scriptsToLoad.push('lunr.multi.js'); + } + if (lang.includes("ja") || lang.includes("jp")) { + scriptsToLoad.push('tinyseg.js'); + } + for (var i=0; i < lang.length; i++) { + if (lang[i] != 'en') { + scriptsToLoad.push(['lunr', lang[i], 'js'].join('.')); + } + } + } + loadScripts(scriptsToLoad, onScriptsLoaded); +} + +function onScriptsLoaded () { + console.log('All search scripts loaded, building Lunr index...'); + if (data.config && data.config.separator && data.config.separator.length) { + lunr.tokenizer.separator = new RegExp(data.config.separator); + } + + if (data.index) { + index = lunr.Index.load(data.index); + data.docs.forEach(function (doc) { + documents[doc.location] = doc; + }); + console.log('Lunr pre-built index loaded, search ready'); + } else { + index = lunr(function () { + if (lang.length === 1 && lang[0] !== "en" && lunr[lang[0]]) { + this.use(lunr[lang[0]]); + } else if (lang.length > 1) { + this.use(lunr.multiLanguage.apply(null, lang)); // spread operator not supported in all browsers: https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/Spread_operator#Browser_compatibility + } + this.field('title'); + this.field('text'); + this.ref('location'); + + for (var i=0; i < data.docs.length; i++) { + var doc = data.docs[i]; + this.add(doc); + documents[doc.location] = doc; + } + }); + console.log('Lunr index built, search ready'); + } + allowSearch = true; + postMessage({config: data.config}); + postMessage({allowSearch: allowSearch}); +} + +function init () { + var oReq = new XMLHttpRequest(); + oReq.addEventListener("load", onJSONLoaded); + var index_path = base_path + '/search_index.json'; + if( 'function' === typeof importScripts ){ + index_path = 'search_index.json'; + } + oReq.open("GET", index_path); + oReq.send(); +} + +function search (query) { + if (!allowSearch) { + console.error('Assets for search still loading'); + return; + } + + var resultDocuments = []; + var results = index.search(query); + for (var i=0; i < results.length; i++){ + var result = results[i]; + doc = documents[result.ref]; + doc.summary = doc.text.substring(0, 200); + resultDocuments.push(doc); + } + return resultDocuments; +} + +if( 'function' === typeof importScripts ) { + onmessage = function (e) { + if (e.data.init) { + init(); + } else if (e.data.query) { + postMessage({ results: search(e.data.query) }); + } else { + console.error("Worker - Unrecognized message: " + e); + } + }; +} diff --git a/sitemap.xml b/sitemap.xml new file mode 100644 index 0000000..746c1aa --- /dev/null +++ b/sitemap.xml @@ -0,0 +1,13 @@ + +