FrEIA.modules package#

Subclasses of torch.nn.Module, that are reversible and can be used in the nodes of the GraphINN class. The only additional things that are needed compared to the base class is an @staticmethod otuput_dims, and the ‘rev’-argument of the forward-method.

Abstract template:

InvertibleModule

Coupling blocks:

AllInOneBlock
NICECouplingBlock
RNVPCouplingBlock
GLOWCouplingBlock
GINCouplingBlock
AffineCouplingOneSided
ConditionalAffineTransform
RationalQuadraticSpline

Reshaping:

IRevNetDownsampling
IRevNetUpsampling
HaarDownsampling
HaarUpsampling
Flatten
Reshape

Graph topology:

Split
Concat

Other learned transforms:

ActNorm
IResNetLayer
InvAutoAct
InvAutoActFixed
InvAutoActTwoSided
InvAutoConv2D
InvAutoFC
LearnedElementwiseScaling
OrthogonalTransform
HouseholderPerm
ElementwiseRationalQuadraticSpline

Fixed (non-learned) transforms:

PermuteRandom
FixedLinearTransform
Fixed1x1Conv
InvertibleSigmoid

Abstract template#

class FrEIA.modules.InvertibleModule(dims_in: List[Tuple[int]], dims_c: List[Tuple[int]] | None = None)[source]#

Base class for all invertible modules in FrEIA.

Given module, an instance of some InvertibleModule. This module shall be invertible in its input dimensions, so that the input can be recovered by applying the module in backwards mode (rev=True), not to be confused with pytorch.backward() which computes the gradient of an operation:

x = torch.randn(BATCH_SIZE, DIM_COUNT)
c = torch.randn(BATCH_SIZE, CONDITION_DIM)

# Forward mode
z, jac = module([x], [c], jac=True)

# Backward mode
x_rev, jac_rev = module(z, [c], rev=True)

The module returns \(\log \det J = \log \left| \det \frac{\partial f}{\partial x} \right|\) of the operation in forward mode, and \(-\log | \det J | = \log \left| \det \frac{\partial f^{-1}}{\partial z} \right| = -\log \left| \det \frac{\partial f}{\partial x} \right|\) in backward mode (rev=True).

Then, torch.allclose(x, x_rev) == True and torch.allclose(jac, -jac_rev) == True.

__init__(dims_in: List[Tuple[int]], dims_c: List[Tuple[int]] | None = None)[source]#

Parameters:

dims_in – list of tuples specifying the shape of the inputs to this operator: dims_in = [shape_x_0, shape_x_1, ...]
dims_c – list of tuples specifying the shape of the conditions to this operator.

forward(x_or_z: Iterable[Tensor], c: Iterable[Tensor] | None = None, rev: bool = False, jac: bool = True) → Tuple[Tuple[Tensor], Tensor][source]#

Perform a forward (default, rev=False) or backward pass (rev=True) through this module/operator.

Note to implementers:

Subclasses MUST return a Jacobian when jac=True, but CAN return a valid Jacobian when jac=False (not punished). The latter is only recommended if the computation of the Jacobian is trivial.
Subclasses MUST follow the convention that the returned Jacobian be consistent with the evaluation direction. Let’s make this more precise: Let \(f\) be the function that the subclass represents. Then:

\[\begin{split}J &= \log \det \frac{\partial f}{\partial x} \\ -J &= \log \det \frac{\partial f^{-1}}{\partial z}.\end{split}\]

Any subclass MUST return \(J\) for forward evaluation (rev=False), and \(-J\) for backward evaluation (rev=True).

Parameters:

x_or_z – input data (array-like of one or more tensors)
c – conditioning data (array-like of none or more tensors)
rev – perform backward pass
jac – return Jacobian associated to the direction

log_jacobian(*args, **kwargs)[source]#: This method is deprecated, and does nothing except raise a warning.

output_dims(input_dims: List[Tuple[int]]) → List[Tuple[int]][source]#

Used for shape inference during construction of the graph. MUST be implemented for each subclass of InvertibleModule.

Parameters:: input_dims – A list with one entry for each input to the module. Even if the module only has one input, must be a list with one entry. Each entry is a tuple giving the shape of that input, excluding the batch dimension. For example for a module with one input, which receives a 32x32 pixel RGB image, input_dims would be [(3, 32, 32)]
Returns:: A list structured in the same way as input_dims. Each entry represents one output of the module, and the entry is a tuple giving the shape of that output. For example if the module splits the image into a right and a left half, the return value should be [(3, 16, 32), (3, 16, 32)]. It is up to the implementor of the subclass to ensure that the total number of elements in all inputs and all outputs is consistent.

Coupling blocks#

class FrEIA.modules.AllInOneBlock(dims_in, dims_c=[], subnet_constructor: Callable | None = None, affine_clamping: float = 2.0, gin_block: bool = False, global_affine_init: float = 1.0, global_affine_type: str = 'SOFTPLUS', permute_soft: bool = False, learned_householder_permutation: int = 0, reverse_permutation: bool = False)[source]#

Module combining the most common operations in a normalizing flow or similar model.

It combines affine coupling, permutation, and global affine transformation (‘ActNorm’). It can also be used as GIN coupling block, perform learned householder permutations, and use an inverted pre-permutation. The affine transformation includes a soft clamping mechanism, first used in Real-NVP. The block as a whole performs the following computation:

\[y = V\,R \; \Psi(s_\mathrm{global}) \odot \mathrm{Coupling}\Big(R^{-1} V^{-1} x\Big)+ t_\mathrm{global}\]

The inverse pre-permutation of x (i.e. \(R^{-1} V^{-1}\)) is optional (see reverse_permutation below).
The learned householder reflection matrix \(V\) is also optional all together (see learned_householder_permutation below).
For the coupling, the input is split into \(x_1, x_2\) along the channel dimension. Then the output of the coupling operation is the two halves \(u = \mathrm{concat}(u_1, u_2)\).

\[\begin{split}u_1 &= x_1 \odot \exp \Big( \alpha \; \mathrm{tanh}\big( s(x_2) \big)\Big) + t(x_2) \\ u_2 &= x_2\end{split}\]

Because \(\mathrm{tanh}(s) \in [-1, 1]\), this clamping mechanism prevents exploding values in the exponential. The hyperparameter \(\alpha\) can be adjusted.

__init__(dims_in, dims_c=[], subnet_constructor: Callable | None = None, affine_clamping: float = 2.0, gin_block: bool = False, global_affine_init: float = 1.0, global_affine_type: str = 'SOFTPLUS', permute_soft: bool = False, learned_householder_permutation: int = 0, reverse_permutation: bool = False)[source]#

Parameters:

subnet_constructor – class or callable f, called as f(channels_in, channels_out) and should return a torch.nn.Module. Predicts coupling coefficients \(s, t\).
affine_clamping – clamp the output of the multiplicative coefficients before exponentiation to +/- affine_clamping (see \(\alpha\) above).
gin_block – Turn the block into a GIN block from Sorrenson et al, 2019. Makes it so that the coupling operations as a whole is volume preserving.
global_affine_init – Initial value for the global affine scaling \(s_\mathrm{global}\).
global_affine_init – 'SIGMOID', 'SOFTPLUS', or 'EXP'. Defines the activation to be used on the beta for the global affine scaling (\(\Psi\) above).
permute_soft – bool, whether to sample the permutation matrix \(R\) from \(SO(N)\), or to use hard permutations instead. Note, permute_soft=True is very slow when working with >512 dimensions.
learned_householder_permutation – Int, if >0, turn on the matrix \(V\) above, that represents multiple learned householder reflections. Slow if large number. Dubious whether it actually helps network performance.
reverse_permutation – Reverse the permutation before the block, as introduced by Putzky et al, 2019. Turns on the \(R^{-1} V^{-1}\) pre-multiplication above.

class FrEIA.modules.NICECouplingBlock(dims_in, dims_c=[], subnet_constructor: callable | None = None, split_len: float | int = 0.5)[source]#

Coupling Block following the NICE (Dinh et al, 2015) design. The inputs are split in two halves. For 2D, 3D, 4D inputs, the split is performed along the channel dimension. Then, residual coefficients are predicted by two subnetworks that are added to each half in turn.

__init__(dims_in, dims_c=[], subnet_constructor: callable | None = None, split_len: float | int = 0.5)[source]#