helios.nn.utils

Attributes

NETWORK_REGISTRY

Global instance of the registry for networks.

Functions

create_network(→ torch.nn.Module)	Create the network for the given type.
default_init_weights(→ None)	Initialize network weights.

Module Contents

helios.nn.utils.NETWORK_REGISTRY

Global instance of the registry for networks.

Example

import helios.nn as hln

# This automatically registers your network.
@hln.NETWORK_REGISTRY.register
class MyNetwork:
    ...

# Alternatively you can manually register a network like this:
hln.NETWORK_REGISTRY.register(MyNetwork)

helios.nn.utils.create_network(type_name: str, *args: Any, **kwargs: Any) → torch.nn.Module

Create the network for the given type.

Parameters:

• type_name – the type of the network to create.

• args – positional arguments to pass into the network.

• kwargs – keyword arguments to pass into the network.

Returns:

The network.

helios.nn.utils.default_init_weights(module_list: list[torch.nn.Module] | torch.nn.Module, scale: float = 1, bias_fill: float = 0, **kwargs: Any) → None

Initialize network weights.

Specifically, this function will default initialize the following types of blocks:

• torch.nn.Conv2d,

• torch.nn.Linear,

• torch.nn.modules.batchnorm._BatchNorm

Parameters:

• module_list – the list of modules to initialize.

• scale – scale initialized weights, especially for residual blocks. Defaults to 1.

• bias_fill – bias fill value. Defaults to 0.

• kwargs – keyword arguments for the torch.nn.init.kaiming_normal_ function.