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)

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)

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]]]

Instantiates the module and determines the output dimension by calling InvertibleModule#output_dims.

class FrEIA.framework.GraphINN(node_list, force_tuple_output=False, verbose=False)

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)

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]

Forward or backward computation of the whole net.

get_module_by_name(name) → Optional[torch.nn.modules.module.Module]

Return module of the first node in the graph with the provided name.

get_node_by_name(name) → Optional[FrEIA.framework.graph_inn.Node]

Return the first node in the graph with the provided name.

log_jacobian_numerical(x, c=None, rev=False, h=0.0001)

Approximate log Jacobian determinant via finite differences.

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

training: bool

class FrEIA.framework.InputNode(*dims: int, name=None)

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)

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]]]

Instantiates the module and determines the output dimension by calling InvertibleModule#output_dims.

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)

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)

Initialize self. See help(type(self)) for accurate signature.

__module__ = 'FrEIA.framework.graph_inn'

__repr__()

Return repr(self).

__str__()

Return str(self).

__weakref__

list of weak references to the object (if defined)

build_module(condition_shapes, input_shapes) → Tuple[FrEIA.modules.base.InvertibleModule, List[Tuple[int]]]

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]]

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)

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)

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]]]

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)

Bases: FrEIA.framework.graph_inn.GraphINN

__init__(node_list, ind_in=None, ind_out=None, verbose=True, force_tuple_output=False)

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]

Forward or backward computation of the whole net.

training: bool

class FrEIA.framework.ReversibleSequential(*dims: int)

Bases: FrEIA.framework.sequence_inn.SequenceINN

__init__(*dims: int)

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)

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)

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)

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]

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)

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

training: bool

Welcome to FrEIA’s documentation! FrEIA.modules package