Top

FrEIA module

Framework for Easily Invertible Architectures. Module to construct invertible networks with pytorch, based on a graph structure of operations.

Show source ≡

'''Framework for Easily Invertible Architectures.
Module to construct invertible networks with pytorch, based on a graph
structure of operations.'''
from . import framework
from . import modules
from . import dummy_modules

__all__ = ["framework", "modules"]

Sub-modules

FrEIA.framework

The framework module contains the logic used in building the graph and inferring the order that the nodes have to be executed in forward and backward direction.

FrEIA.modules

Subclasses of torch.nn.Module, that are reversible and can be used in the nodes of the ReversibleGraphNet 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.

Coupling blocks:

  • NICECouplingBlock...