:py:mod:`anomalib.models.components.freia.modules` ================================================== .. py:module:: anomalib.models.components.freia.modules .. autoapi-nested-parse:: Modules. Submodules ---------- .. toctree:: :titlesonly: :maxdepth: 1 all_in_one_block/index.rst base/index.rst Package Contents ---------------- Classes ~~~~~~~ .. autoapisummary:: anomalib.models.components.freia.modules.AllInOneBlock anomalib.models.components.freia.modules.InvertibleModule .. py:class:: 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) Bases: :py:obj:`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. :math:`R^{-1} V^{-1}`) is optional (see ``reverse_permutation`` below). - The learned householder reflection matrix :math:`V` is also optional all together (see ``learned_householder_permutation`` below). - For the coupling, the input is split into :math:`x_1, x_2` along the channel dimension. Then the output of the coupling operation is the two halves :math:`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 :math:`\\mathrm{tanh}(s) \\in [-1, 1]`, this clamping mechanism prevents exploding values in the exponential. The hyperparameter :math:`\\alpha` can be adjusted. .. py:method:: _construct_householder_permutation() Compute a permutation matrix. Compute a permutation matrix from the reflection vectors that are learned internally as nn.Parameters. .. py:method:: _permute(x, rev=False) Perform permutation. Performs the permutation and scaling after the coupling operation. Returns transformed outputs and the LogJacDet of the scaling operation. .. py:method:: _pre_permute(x, rev=False) Permute before the coupling block, only used if reverse_permutation is set. .. py:method:: _affine(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. .. py:method:: forward(x, c=[], rev=False, jac=True) See base class docstring. .. py:method:: output_dims(input_dims) Output Dims. .. py:class:: InvertibleModule(dims_in: Iterable[Tuple[int]], dims_c: Iterable[Tuple[int]] = None) Bases: :py:obj:`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 :math:`\\log \\det J = \\log \\left| \\det \\frac{\\partial f}{\\partial x} \\right|` of the operation in forward mode, and :math:`-\\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``. .. py:method:: forward(x_or_z: Iterable[torch.Tensor], c: Iterable[torch.Tensor] = None, rev: bool = False, jac: bool = True) -> Tuple[Tuple[torch.Tensor], torch.Tensor] :abstractmethod: 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 :math:`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 :math:`J` for forward evaluation (``rev=False``), and :math:`-J` for backward evaluation (``rev=True``). :param x_or_z: input data (array-like of one or more tensors) :param c: conditioning data (array-like of none or more tensors) :param rev: perform backward pass :param jac: return Jacobian associated to the direction .. py:method:: log_jacobian(*args, **kwargs) This method is deprecated, and does nothing except raise a warning. .. py:method:: output_dims(input_dims: List[Tuple[int]]) -> List[Tuple[int]] :abstractmethod: Use for shape inference during construction of the graph. MUST be implemented for each subclass of ``InvertibleModule``. :param 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.