FrEIA.framework package¶
Module contents¶
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.
-
class
FrEIA.framework.
ConditionNode
(*dims: int, name=None)[source]¶ Bases:
FrEIA.framework.graph_inn.Node
Special type of node that represents contitional input to the internal networks inside coupling layers.
-
__init__
(*dims: int, name=None)[source]¶ Initialize self. See help(type(self)) for accurate signature.
-
__module__
= 'FrEIA.framework.graph_inn'¶
-
-
class
FrEIA.framework.
GraphINN
(node_list, force_tuple_output=False, verbose=False)[source]¶ Bases:
FrEIA.modules.base.InvertibleModule
This class represents the invertible net itself. It is a subclass of InvertibleModule and supports the same methods.
The forward method has an additional option ‘rev’, with which the net can be computed in reverse. Passing jac to the forward method additionally computes the log determinant of the (inverse) Jacobian of the forward (backward) pass.
-
__init__
(node_list, force_tuple_output=False, verbose=False)[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.
-
__module__
= 'FrEIA.framework.graph_inn'¶
-
forward
(x_or_z: Union[torch.Tensor, Iterable[torch.Tensor]], c: Optional[Iterable[torch.Tensor]] = None, rev: bool = False, jac: bool = True, intermediate_outputs: bool = False, x: None = None) → Tuple[Tuple[torch.Tensor], torch.Tensor][source]¶ Forward or backward computation of the whole net.
-
get_module_by_name
(name) → Optional[torch.nn.modules.module.Module][source]¶ Return module of the first node in the graph with the provided name.
-
get_node_by_name
(name) → Optional[FrEIA.framework.graph_inn.Node][source]¶ Return the first node in the graph with the provided name.
-
log_jacobian_numerical
(x, c=None, rev=False, h=0.0001)[source]¶ Approximate log Jacobian determinant via finite differences.
-
training
: bool¶
-
-
class
FrEIA.framework.
InputNode
(*dims: int, name=None)[source]¶ Bases:
FrEIA.framework.graph_inn.Node
Special type of node that represents the input data of the whole net (or the output when running reverse)
-
__init__
(*dims: int, name=None)[source]¶ Initialize self. See help(type(self)) for accurate signature.
-
__module__
= 'FrEIA.framework.graph_inn'¶
-
-
class
FrEIA.framework.
Node
(inputs: Union[FrEIA.framework.graph_inn.Node, Tuple[FrEIA.framework.graph_inn.Node, int], Iterable[Tuple[FrEIA.framework.graph_inn.Node, int]]], module_type, module_args: dict, conditions=None, name=None)[source]¶ Bases:
object
The Node class represents one transformation in the graph, with an arbitrary number of in- and outputs.
The user specifies the input, and the underlying module computes the number of outputs.
-
__dict__
= mappingproxy({'__module__': 'FrEIA.framework.graph_inn', '__doc__': '\n The Node class represents one transformation in the graph, with an\n arbitrary number of in- and outputs.\n\n The user specifies the input, and the underlying module computes the\n number of outputs.\n ', '__init__': <function Node.__init__>, 'build_module': <function Node.build_module>, 'parse_inputs': <function Node.parse_inputs>, '__str__': <function Node.__str__>, '__repr__': <function Node.__repr__>, '__dict__': <attribute '__dict__' of 'Node' objects>, '__weakref__': <attribute '__weakref__' of 'Node' objects>, '__annotations__': {'outputs': 'List[Tuple[Node, int]]'}})¶
-
__init__
(inputs: Union[FrEIA.framework.graph_inn.Node, Tuple[FrEIA.framework.graph_inn.Node, int], Iterable[Tuple[FrEIA.framework.graph_inn.Node, int]]], module_type, module_args: dict, conditions=None, name=None)[source]¶ Initialize self. See help(type(self)) for accurate signature.
-
__module__
= 'FrEIA.framework.graph_inn'¶
-
__weakref__
¶ list of weak references to the object (if defined)
-
build_module
(condition_shapes, input_shapes) → Tuple[FrEIA.modules.base.InvertibleModule, List[Tuple[int]]][source]¶ Instantiates the module and determines the output dimension by calling InvertibleModule#output_dims.
-
parse_inputs
(inputs: Union[FrEIA.framework.graph_inn.Node, Tuple[FrEIA.framework.graph_inn.Node, int], Iterable[Tuple[FrEIA.framework.graph_inn.Node, int]]]) → List[Tuple[FrEIA.framework.graph_inn.Node, int]][source]¶ Converts specified inputs to a node to a canonical format. Inputs can be specified in three forms:
a single node, then this nodes first output is taken as input
a single tuple (node, idx), specifying output idx of node
a list of tuples [(node, idx)], each specifying output idx of node
All such formats are converted to the last format.
-
-
class
FrEIA.framework.
OutputNode
(in_node: Union[FrEIA.framework.graph_inn.Node, Tuple[FrEIA.framework.graph_inn.Node, int]], name=None)[source]¶ Bases:
FrEIA.framework.graph_inn.Node
Special type of node that represents the output of the whole net (or the input when running in reverse).
-
__init__
(in_node: Union[FrEIA.framework.graph_inn.Node, Tuple[FrEIA.framework.graph_inn.Node, int]], name=None)[source]¶ Initialize self. See help(type(self)) for accurate signature.
-
__module__
= 'FrEIA.framework.graph_inn'¶
-
build_module
(condition_shapes, input_shapes) → Tuple[None, List[Tuple[int]]][source]¶ Instantiates the module and determines the output dimension by calling InvertibleModule#output_dims.
-
outputs
: List[Tuple[FrEIA.framework.graph_inn.Node, int]]¶
-
-
class
FrEIA.framework.
ReversibleGraphNet
(node_list, ind_in=None, ind_out=None, verbose=True, force_tuple_output=False)[source]¶ Bases:
FrEIA.framework.graph_inn.GraphINN
-
__init__
(node_list, ind_in=None, ind_out=None, verbose=True, force_tuple_output=False)[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.
-
__module__
= 'FrEIA.framework.reversible_graph_net'¶
-
forward
(x_or_z: Union[torch.Tensor, Iterable[torch.Tensor]], c: Optional[Iterable[torch.Tensor]] = None, rev: bool = False, jac: bool = True, intermediate_outputs: bool = False) → Tuple[Tuple[torch.Tensor], torch.Tensor][source]¶ Forward or backward computation of the whole net.
-
training
: bool¶
-
-
class
FrEIA.framework.
ReversibleSequential
(*dims: int)[source]¶ Bases:
FrEIA.framework.sequence_inn.SequenceINN
-
__init__
(*dims: int)[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.
-
__module__
= 'FrEIA.framework.reversible_sequential_net'¶
-
training
: bool¶
-
-
class
FrEIA.framework.
SequenceINN
(*dims: int, force_tuple_output=False)[source]¶ Bases:
FrEIA.modules.base.InvertibleModule
Simpler than FrEIA.framework.GraphINN: Only supports a sequential series of modules (no splitting, merging, branching off). Has an append() method, to add new blocks in a more simple way than the computation-graph based approach of GraphINN. For example:
``` inn = SequenceINN(channels, dims_H, dims_W)
- for i in range(n_blocks):
inn.append(FrEIA.modules.AllInOneBlock, clamp=2.0, permute_soft=True)
inn.append(FrEIA.modules.HaarDownsampling) # and so on ```
-
__init__
(*dims: int, force_tuple_output=False)[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.
-
__module__
= 'FrEIA.framework.sequence_inn'¶
-
append
(module_class, cond=None, cond_shape=None, **kwargs)[source]¶ Append a reversible block from FrEIA.modules to the network. module_class: Class from FrEIA.modules. cond (int): index of which condition to use (conditions will be passed as list to forward()).
Conditioning nodes are not needed for SequenceINN.
cond_shape (tuple[int]): the shape of the condition tensor. **kwargs: Further keyword arguments that are passed to the constructor of module_class (see example).
-
forward
(x_or_z: torch.Tensor, c: Optional[Iterable[torch.Tensor]] = None, rev: bool = False, jac: bool = True) → Tuple[torch.Tensor, torch.Tensor][source]¶ Executes the sequential INN in forward or inverse (rev=True) direction.
- Parameters
x_or_z – input tensor (in contrast to GraphINN, a list of tensors is not supported, as SequenceINN only has one input).
c – list of conditions.
rev – whether to compute the network forward or reversed.
jac – whether to compute the log jacobian
- Returns
network output. jac (Tensor): log-jacobian-determinant.
- Return type
z_or_x (Tensor)
-
training
: bool¶