Network#

Classes#

class BaseNetwork[source]#

Bases: Module, ABC, Generic[ModuleListT]

Base class for all networks in NetLoader

Attributes:

namestr: Name of the network, used for saving
versionstr: NetLoader version string
checkpointslist[TensorListLike]: Outputs from each checkpoint with each Tensor having shape (N,…) and type float, where N is the batch size
kl_lossTensor: KL divergence loss on the latent space of shape (1) and type float, if using a sample layer
layersModuleListT: Network construction

Initialize internal Module state, shared by both nn.Module and ScriptModule.

abstractmethod forward(x)[source]#

Forward pass of the network.

Parameters:

xTensorListLike: Input tensor with shape (N,…) and type float, where N is the batch size

Returns:

TensorListLike: Output tensor from the network with shape (N,…) and type float

get_config(dict_=True)[source]#

Returns the network configuration.

Overloads:

get_config(self, dict_: Literal[True] = ...) → dict[str, Any]

get_config(self, dict_: Literal[False] = ...) → Config

Parameters:

dict_bool, Optional: If configuration should be returned as a dictionary, default = True

Returns:

dict[str, Any] | Config: Network configuration

to(*args, **kwargs)[source]#: See torch.nn.Module.to().

class CompatibleNetwork(net, *, name='')[source]#

Bases: BaseNetwork[Module | ModuleList]

A wrapper for Module that ensures compatibility with BaseArchitecture by adding required attributes.

Attributes:

namestr: Name of the network, used for saving
versionstr: NetLoader version string
checkpointslist[TensorListLike]: Outputs from each checkpoint with each Tensor having shape (N,…) and type float, where N is the batch size
kl_lossTensor: KL divergence loss on the latent space of shape (1) and type float, if using a sample layer
layersModule | ModuleList: Network construction

Parameters:

netModule | ModuleList: The neural network module to wrap or list of layers in the network
namestr, Optional: Name of the network, used for saving

forward(x)[source]#

Forward pass of the network.

Parameters:

xTensorListLike: Input tensor with shape (N,…) and type float, where N is the batch size

Returns:

TensorListLike: Output tensor from the network with shape (N,…) and type float

to(*args, **kwargs)[source]#: See torch.nn.Module.to().

class Network(name, config, in_shape, out_shape, *, suppress_warning=False, root='', defaults=None)[source]#

Bases: BaseNetwork[TypedModuleList[BaseLayer]]

Constructs a neural network from a configuration file.

Attributes:

groupint: Which group is the active group if a layer has the group attribute
layer_numint | NoneType: Number of layers to use, if None use all layers
namestr: Name of the network, used for saving
versionstr: NetLoader version string
checkpointslist[TensorListLike]: Outputs from each checkpoint with each Tensor having shape (N,…) and type float, where N is the batch size
kl_lossTensor: KL divergence loss on the latent space of shape (1) and type float, if using a sample layer
layersTypedModuleList[BaseLayer]: Network layers
shapesShapes: Layer output shapes
check_shapesShapes: Checkpoint output shapes

Parameters:

namestr: Name of the network configuration file
configstr | Config: Path to the network config directory or configuration dictionary
in_shapelist[int] | list[list[int]] | tuple[int, …]: Shape of the input tensor(s), excluding batch size
out_shapelist[int]: Shape of the output tensor, excluding batch size
suppress_warningbool, Optional: If output shape mismatch warning should be suppressed, default = False
rootstr, Optional: Root directory to prepend to config file path
defaultsdict[str, Any] | NoneType, Optional: Default values for the parameters for each type of layer

__getitem__(idx)[source]#

Returns the layer at the specified index.

Overloads:

__getitem__(self, idx: int) → BaseLayer

__getitem__(self, idx: slice) → TypedModuleList[BaseLayer]

Parameters:

idxint | slice: Index or slice of the layer(s) to return

Returns:

Module | TypedModuleList[BaseLayer]: Layer(s) at the specified index or slice

__iter__()[source]#

Returns an iterator over the layers in the network.

Returns:

Iterator[Module]: Iterator over the layers in the network

__len__()[source]#

Returns the number of layers in the network.

Returns:

int: Number of layers in the network

forward(x)[source]#

Forward pass of the network.

Parameters:

xTensorListLike: Input tensor(s) with shape (N,…) and type float, where N is the batch size

Returns:

NormalizingFlow | TensorListLike: Output tensor from the network with shape (N,…) and type float, or NormalizingFlow if the last layer is from layers.flows

get_config(dict_=True)[source]#

Returns the network configuration.

Overloads:

get_config(self, dict_: Literal[True] = ...) → dict[str, Any]

get_config(self, dict_: Literal[False] = ...) → Config

Parameters:

dict_bool, Optional: If configuration should be returned as a dictionary, default = True

Returns:

dict[str, Any] | Config: Network configuration

to(*args, **kwargs)[source]#: See torch.nn.Module.to().

class Branch(net_check, branches, shapes, *, checkpoint=True, dim=0, channels=None, root='', method='concat', shape=None, defaults=None, **kwargs)[source]#

Bases: BaseLayer

Creates a branching layer that processes the input through multiple branches and combines the outputs.

Attributes:

groupint: Layer group, if 0 it will always be used, else it will only be used if its group matches the Networks
descriptionstr: Description of the layer
brancheslist[BaseLayer]: List of subnetworks for each branch
headConcatenate | Pack | Shortcut: Layer to combine the outputs from each branch, either concatenation, packing or summation

Parameters:

net_checkbool: If layer index should be relative to checkpoint layers
brancheslist[list[dict[str, Any]]]: List of layer configurations for each branch
shapesShapes: Shape of the outputs from each layer
checkpointbool, Optional: If layer index should be relative to checkpoint layers, default = True
dimint, Optional: Dimension to concatenate along if method is ‘concat’, default = 0
channelsint, Optional: Number of output channels, won’t be used if shape is provided, if channels and shape aren’t provided, the input dimensions will be preserved
rootstr, Optional: Root directory to prepend to config file path
methodLiteral[‘concat’, ‘pack’, ‘sum’], Optional: Method to combine the outputs from each branch, either concatenation, packing or summation, default = ‘concat’
shapelist[int] | NoneType, Optional: Output shape of the block, will be used if provided; otherwise, channels will be used
defaultsdict[str, Any] | NoneType, Optional: Default values for the parameters for each type of layer
**kwargs: Leftover parameters to pass to base layer for checking

extra_repr()[source]#

Displays layer parameters when printing the network.

Returns:

str: Layer parameters

forward(x, *_, **__)[source]#

Forward pass for the Branch layer.

Parameters:

xTensorListT: Input tensor(s) with dtype float32 and shape (N,…), where N is batch size

Returns:

TensorListT: Output tensor(s) from the subnetwork with dtype float32 and shape (N,…)

to(*args, **kwargs)[source]#: See torch.nn.Module.to().

class Composite(net_check, name, shapes, *, checkpoint=True, channels=None, root='', config_dir='', shape=None, defaults=None, config=None, **kwargs)[source]#

Bases: BaseLayer

Creates a subnetwork from a configuration file.

Attributes:

groupint: Layer group, if 0 it will always be used, else it will only be used if its group matches the Networks
descriptionstr: Description of the layer
netNetwork: Subnetwork for the Composite layer

Parameters:

net_checkbool: If layer index should be relative to checkpoint layers
namestr: Name of the subnetwork
shapesShapes: Shape of the outputs from each layer
checkpointbool, Optional: If layer index should be relative to checkpoint layers, default = True
channelsint, Optional: Number of output channels, won’t be used if shape is provided, if channels and shape aren’t provided, the input dimensions will be preserved
rootstr, Optional: Root directory to prepend to config file path
config_dirstr, Optional: Path to the directory with the network configuration file, won’t be used if config is provided
shapelist[int] | NoneType, Optional: Output shape of the block, will be used if provided; otherwise, channels will be used
defaultsdict[str, Any] | NoneType, Optional: Default values for the parameters for each type of layer
configdict[str, Any] | Config | NoneType, Optional: Network configuration dictionary
**kwargs: Leftover parameters to pass to base layer for checking

extra_repr()[source]#

Displays layer parameters when printing the network.

Returns:

str: Layer parameters

forward(x, *_, **__)[source]#

Forward pass for the Composite layer.

Parameters:

xTensorListT: Input tensor(s) with dtype float32 and shape (N,…), where N is batch size

Returns:

TensorListT: Output tensor(s) from the subnetwork with dtype float32 and shape (N,…)

to(*args, **kwargs)[source]#: See torch.nn.Module.to().