anomalib.models.components.freia.modules

Modules.

Submodules

Package Contents

Classes

AllInOneBlock

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

InvertibleModule

Base class for all invertible modules in FrEIA.
class anomalib.models.components.freia.modules.AllInOneBlock(dims_in, dims_c=[], subnet_constructor: Callable = 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]

Bases: anomalib.models.components.freia.modules.base.InvertibleModule

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: .. math:

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)\). .. math:

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

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

_construct_householder_permutation(self)

Compute a permutation matrix.

Compute a permutation matrix from the reflection vectors that are learned internally as nn.Parameters.

_permute(self, x, rev=False)

Perform permutation.

Performs the permutation and scaling after the coupling operation. Returns transformed outputs and the LogJacDet of the scaling operation.

_pre_permute(self, x, rev=False)

Permute before the coupling block, only used if reverse_permutation is set.

_affine(self, x, a, rev=False)

Perform affine coupling operation.

Given the passive half, and the pre-activation outputs of the coupling subnetwork, perform the affine coupling operation. Returns both the transformed inputs and the LogJacDet.

forward(self, x, c=[], rev=False, jac=True)

See base class docstring.

output_dims(self, input_dims)

Output Dims.

class anomalib.models.components.freia.modules.InvertibleModule(dims_in: Iterable[Tuple[int]], dims_c: Iterable[Tuple[int]] = None)[source]

Bases: torch.nn.Module

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.

abstract forward(self, x_or_z: Iterable[torch.Tensor], c: Iterable[torch.Tensor] = None, rev: bool = False, jac: bool = True) → Tuple[Tuple[torch.Tensor], torch.Tensor]

Forward/Backward Pass.

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: .. math:

    J &= \\log \\det \\frac{\\partial f}{\\partial x} \\\\
-J &= \\log \\det \\frac{\\partial f^{-1}}{\\partial z}.

    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(self, *args, **kwargs)

This method is deprecated, and does nothing except raise a warning.

abstract output_dims(self, input_dims: List[Tuple[int]]) → List[Tuple[int]]

Use 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.