konfai.network package

Submodules

konfai.network.blocks module

Reusable network blocks and tensor graph helpers for KonfAI models.

class konfai.network.blocks.NormMode(*values)

Bases: Enum

Enumeration of normalization layers supported by KonfAI blocks.

NONE = (0,)

BATCH = 1

INSTANCE = 2

GROUP = 3

LAYER = 4

SYNCBATCH = 5

INSTANCE_AFFINE = 6

konfai.network.blocks.get_norm(norm_mode, channels, dim)

Instantiate the normalization layer matching the requested mode.

Return type:

Module | None

class konfai.network.blocks.UpsampleMode(*values)

Bases: Enum

CONV_TRANSPOSE = (0,)

UPSAMPLE = (1,)

class konfai.network.blocks.DownsampleMode(*values)

Bases: Enum

MAXPOOL = (0,)

AVGPOOL = (1,)

CONV_STRIDE = 2

konfai.network.blocks.get_torch_module(name_fonction, dim=None)

Return a dimensional PyTorch module class such as Conv2d or Conv3d.

Return type:

Module

class konfai.network.blocks.BlockConfig(kernel_size=3, stride=1, padding=1, bias=True, activation='ReLU', norm_mode='NONE')

Bases: object

Configuration object describing one convolutional block stage.

norm: NormMode | Callable[[int], Module] | None

get_conv(in_channels, out_channels, dim)

Return type:

Conv3d

get_norm(channels, dim)

Return type:

Module

get_activation()

Return type:

Module

class konfai.network.blocks.ConvBlock(in_channels, out_channels, block_configs, dim, alias=[[], [], []])

Bases: ModuleArgsDict

Sequential convolution, normalization, and activation block.

class konfai.network.blocks.ResBlock(in_channels, out_channels, block_configs, dim, alias=[[], [], [], [], []])

Bases: ModuleArgsDict

Residual block with optional projection on the skip path.

class konfai.network.blocks.ResidualBlockD(in_channels, out_channels, dim, kernel_size=3, stride=1, conv_bias=True, negative_slope=0.01, norm_mode='INSTANCE_AFFINE')

Bases: ModuleArgsDict

nnU-Net ResNet-D basic residual block, reproduced as a routed KonfAI graph.

Module-for-module and in forward-execution order equal to dynamic_network_architectures.building_blocks.residual.BasicBlockD (the block the nnU-Net ResidualEncoderUNet stacks). It reproduces the ResNet-D structure of He et al., Bag of Tricks for Image Classification with CNNs (CVPR 2019):

• Skip path first (matching BasicBlockD.forward, which evaluates self.skip(x) before the main path): when the block is strided, the residual is downsampled with an AvgPool of kernel = stride (never a strided conv); when the channel count changes, it is then projected with a 1x1 conv whose bias is always False – independent of conv_bias – followed by a norm. When neither applies the skip is the identity (no extra module, the residual falls back to the block input branch).

• Main path on branch 0: Conv(stride) -> Norm -> LeakyReLU -> Conv(stride 1) -> Norm (the second conv carries no activation, exactly as BasicBlockD).

• The two paths are summed and a final LeakyReLU is applied.

Executing the skip modules before the main-path convs is what makes the block transparent to the execution-order weight bridge (konfai.utils.pretrained): its weighted leaves fire in the same order as BasicBlockD (skip conv/norm, then conv1, then conv2).

class konfai.network.blocks.ResNetBasicBlock(in_channels, out_channels, dim, stride=1, conv_bias=False, downsample=None, norm_mode='BATCH')

Bases: ModuleArgsDict

torchvision ResNet BasicBlock reproduced as a routed KonfAI graph (the ResNet-18/34 block).

Module-for-module and in forward-execution order equal to torchvision.models.resnet.BasicBlock (the block the resnet18/resnet34 encoders of segmentation_models_pytorch stack). It reproduces the post-activation residual block of He et al., Deep Residual Learning (CVPR 2016):

• Main path first (matching BasicBlock.forward, which computes the residual branch before it evaluates self.downsample(x)): on branch 1 Conv(stride) -> Norm -> ReLU -> Conv(stride 1) -> Norm (the second conv carries no activation), keeping the block input untouched on branch 0.

• Projection skip on branch 2, evaluated after the main path (exactly as BasicBlock calls self.downsample(x) only after bn2): a 1x1 Conv(stride) -> Norm. It is built when the block is strided or changes channel count; otherwise the residual is the identity (block input).

• The two paths are summed onto branch 0 and a final ReLU is applied.

Evaluating the main-path convs before the downsample projection is what makes the block transparent to the execution-order weight bridge (konfai.utils.pretrained): its weighted leaves fire in the same order as BasicBlock (conv1, bn1, conv2, bn2, downsample.0, downsample.1). conv_bias defaults to False (torchvision/timm ResNets never bias their convs) and the norm is BatchNorm (running stats, so a checkpoint’s BN buffers travel through the bridge and the eval() forward matches).

class konfai.network.blocks.ResidualStage(dim, in_channels, out_channels, n_blocks, stride=1, kernel_size=3, conv_bias=True, negative_slope=0.01, norm_mode='INSTANCE_AFFINE')

Bases: ModuleArgsDict

One encoder resolution stage: a stack of n_blocks ResidualBlockD as a single node.

A generic, reusable building block for residual encoders (e.g. nnU-Net’s ResidualEncoder): the first block carries stride and the in_channels -> out_channels change (nnU-Net strided-conv downsampling), and the remaining n_blocks - 1 blocks are stride 1 at out_channels. The blocks are wired sequentially (each on branch 0), so the whole stage is a drop-in single node with one input and one output whose weighted leaves fire in the same order as the equivalent flat stack – it stays transparent to the execution-order weight bridge (konfai.utils.pretrained).

class konfai.network.blocks.DecoderStage(dim, in_channels, skip_channels, n_conv, kernel_size=3, conv_bias=True, negative_slope=0.01, upsample_stride=2, norm_mode='INSTANCE_AFFINE')

Bases: ModuleArgsDict

One U-Net decoder resolution stage as a single two-input node: upsample, concat skip, conv block.

A generic, reusable building block for U-Net decoders (nnU-Net’s UNetDecoder shares it between the plain and residual encoders). It is a two-input node – in_branch: [coarser, skip] – that runs ConvTranspose(in_channels -> skip_channels, kernel = stride = upsample_stride) on the coarser input, concatenates the encoder skip (transpose output first, then skip), then a ConvBlock of n_conv convs mapping 2 * skip_channels -> skip_channels – each Conv -> InstanceNorm(affine) -> LeakyReLU with same-padding. It yields one output. upsample_stride sets both the transpose kernel and stride so anisotropic plans (per-axis lists) upsample exactly the axes their encoder downsampled.

class konfai.network.blocks.ResNetStage(dim, in_channels, out_channels, n_blocks, stride=1, conv_bias=False, norm_mode='BATCH')

Bases: ModuleArgsDict

One torchvision ResNet encoder stage: a stack of n_blocks ResNetBasicBlock as a single node.

A generic, reusable building block for the ResNet-18/34 encoders that segmentation_models_pytorch stacks under its U-Net / UNet++ decoders (each is torchvision’s ResNet layer1..layer4). The first block carries the stage stride and the in_channels -> out_channels change – its 1x1 projection downsample skip is built automatically when the stride or the channel count changes – and the remaining n_blocks - 1 blocks are stride 1 at out_channels with identity skips. The blocks run sequentially on branch 0, so the whole stage is a drop-in single node with one input and one output whose weighted leaves fire in the same order as the equivalent flat BasicBlock stack: it stays transparent to the execution-order weight bridge (konfai.utils.pretrained). conv_bias defaults to False and the norm to BatchNorm – the torchvision/timm ResNet convention.

class konfai.network.blocks.UNetPlusPlusNode(dim, up_channels, skip_channels, out_channels, n_conv=2, kernel_size=3, conv_bias=False, norm_mode='BATCH')

Bases: ModuleArgsDict

One UNet++ dense-decoder node as a single multi-input node: upsample, dense concat, then Conv-Norm-ReLU.

A generic, reusable building block for the nested (dense) UNet++ decoder of segmentation_models_pytorch (UnetPlusPlusDecoder). It is a multi-input node – in_branch: [coarser, skip_0, skip_1, ...] – that reproduces one grid node x_{d}_{l}:

• Upsample(scale_factor=2, mode='nearest') on the shallower-column predecessor (branch 0);

• a Concat of the upsampled feature FIRST, then every same-resolution dense skip and the matching encoder skip (smp order), when any skip is provided;

• a ConvBlock of n_conv Conv -> BatchNorm -> ReLU blocks (smp’s Conv2dReLU) mapping the concatenated width up_channels + sum(skip_channels) to out_channels.

skip_channels is the list of per-skip channel widths in in_branch order: its length wires the concat (n_skip + 1 inputs) and its sum fixes the conv’s input width, so the node self-describes the dense fusion. An empty skip_channels (the final full-resolution x_0_depth node) drops the concat and convolves the upsampled feature alone. The weighted leaves fire in ConvBlock order, so the node stays transparent to the execution-order weight bridge (konfai.utils.pretrained).

konfai.network.blocks.downsample(in_channels, out_channels, downsample_mode, dim)

Return the downsampling module matching the requested strategy.

Return type:

Module

konfai.network.blocks.upsample(in_channels, out_channels, upsample_mode, dim, kernel_size=2, stride=2)

Return the upsampling module matching the requested strategy.

class konfai.network.blocks.Unsqueeze(dim=0)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class konfai.network.blocks.Permute(dims)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class konfai.network.blocks.ToChannels(dim)

Bases: Permute

class konfai.network.blocks.ToFeatures(dim)

Bases: Permute

class konfai.network.blocks.Add

Bases: Module

forward(*tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Multiply

Bases: Module

forward(*tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.ClipNormalize

Bases: Module

Clip to a stored intensity range, then standardize: (clamp(x, min, max) - mean) / std.

The four scalars are buffers restored from the checkpoint, not config values, so a per-checkpoint input normalization (e.g. a CT window baked into a trained model) travels with its weights. It has no learnable parameters. Used as a declarative model’s first node so that normalization stays part of the model rather than a separate preprocessing step.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Concat

Bases: Module

forward(*tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Print

Bases: Module

Debug block: print the tensor shape and pass it through unchanged.

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Write(path='./Data.mha')

Bases: Module

Debug block: write the first channel of the first sample to disk as a volume.

The destination path is explicit (no silent hardcoded location) and each call warns, since writing to disk inside a forward pass is a development-only side effect.

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Exit

Bases: Module

Debug block: stop the forward pass by raising, to halt at a chosen point.

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Detach

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Negative

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.GetShape

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.ArgMax(dim)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Select(slices)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.NormalNoise(dim=None)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Const(shape, std)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Subset(slices)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.View(size)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.LatentDistribution(shape, latent_dim)

Bases: ModuleArgsDict

class LatentDistributionLinear(shape, latent_dim)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class LatentDistributionDecoder(shape, latent_dim)

Bases: Module

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class LatentDistributionZ

Bases: Module

forward(mu, log_std)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class konfai.network.blocks.Attention(f_g, f_l, f_int, dim)

Bases: ModuleArgsDict

class konfai.network.blocks.PositionalEmbedding(num_tokens, embedding_dim)

Bases: Module

Add a learnable positional embedding to a token sequence.

The input is a token sequence shaped [B, num_tokens, embedding_dim] (batch, sequence, feature), the layout produced by a patch embedding. A single learnable parameter of shape [1, num_tokens, embedding_dim] is broadcast over the batch and added, giving every token position a trainable offset. This is the learnable positional encoding of the Vision Transformer; the parameter is registered directly on the module so it is a self-contained, single-purpose building block.

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

class konfai.network.blocks.MultiHeadSelfAttention(hidden_size, num_heads, qkv_bias=False)

Bases: Module

Multi-head self-attention over a token sequence (the Transformer encoder attention).

Mirrors the self-attention of MONAI’s SABlock/ViT in its default configuration: a single packed qkv linear projects the [B, num_tokens, hidden_size] sequence into queries, keys and values, scaled dot-product attention is computed per head with scale head_dim ** -0.5, and an out_proj linear mixes the concatenated heads back to hidden_size. qkv_bias toggles the bias of the packed projection (the ViT default is False); out_proj always carries a bias. Both projections are child Linear leaves, so the module is transparent to execution-order weight transfer.

forward(tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

konfai.network.network module

Model graph composition, routing, and optimization helpers for KonfAI.

class konfai.network.network.NetState(*values)

Bases: Enum

Execution state of a network inside KonfAI workflows.

TRAIN = (0,)

PREDICTION = 1

class konfai.network.network.PatchIndexed(patch, index)

Bases: object

Track progress while consuming the patches produced by a ModelPatch.

is_full()

Return type:

bool

class konfai.network.network.OptimizerLoader(name='AdamW')

Bases: object

Configuration-aware factory for PyTorch optimizers.

get_optimizer(key, parameter)

Return type:

Optimizer

class konfai.network.network.LRSchedulersLoader(nb_step=0)

Bases: object

Configuration-aware factory for learning-rate schedulers.

getschedulers(key, scheduler_classname, optimizer)

Return type:

_LRScheduler

class konfai.network.network.LossSchedulersLoader(nb_step=0)

Bases: object

Factory for scalar schedulers attached to losses and metrics.

getschedulers(key, scheduler_classname)

Return type:

_LRScheduler

konfai.network.network.build_configured_criterions(criterions_loader, config_key_prefix, configure_attr=None)

Instantiate criteria from a KonfAI config subtree.

Parameters:

• criterions_loader (dict[str, Any]) – Mapping from criterion classpath to criterion attributes.

• config_key_prefix (str) – Config subtree preceding .criterions_loader.<classpath>.

• configure_attr (Callable[[str, Any, Any], None] | None) – Optional callback used to enrich the attribute object with runtime metadata before the criterion is instantiated. It receives the criterion classpath, the attribute object, and the imported module.

Returns:

Instantiated criteria mapped to their configuration attributes.

Return type:

dict[Module, Any]

class konfai.network.network.CriterionsAttr(schedulers={'default|Constant': <konfai.network.network.LossSchedulersLoader object>}, is_loss=True, group=0, start=0, stop=None, accumulation=False)

Bases: object

Metadata describing how a criterion is applied within the model graph.

schedulers: dict[Scheduler, int]

class konfai.network.network.CriterionsLoader(criterions_loader={'default|torch:nn:CrossEntropyLoss|Dice|NCC': <konfai.network.network.CriterionsAttr object>})

Bases: object

Instantiate the criteria attached to one output/target pair.

get_criterions(model_classname, output_group, target_group)

Return type:

dict[Module, CriterionsAttr]

class konfai.network.network.TargetCriterionsLoader(targets_criterions={'Labels': <konfai.network.network.CriterionsLoader object>})

Bases: object

Resolve criteria for all targets associated with one model output.

get_targets_criterions(output_group, model_classname)

Return type:

dict[str, dict[Module, CriterionsAttr]]

class konfai.network.network.Measure(model_classname, outputs_criterions_loader)

Bases: object

Collect, validate, and aggregate losses or metrics across model outputs.

class Loss(name, output_group, target_group, group, is_loss, accumulation)

Bases: object

reset_loss()

Return type:

None

add(weight, value)

Return type:

None

get_last_loss()

Return type:

Tensor

get_loss()

Return type:

Tensor

outputs_criterions: dict[str, dict[str, dict[Module, CriterionsAttr]]]

scaler: GradScaler | None

init(model, group_dest)

Return type:

None

update(output_group, output, batch_data_with_attribute, it, nb_patch, training)

Return type:

None

get_loss()

Return type:

list[Tensor]

reset_loss()

Return type:

None

get_last_values(n=1)

Return type:

dict[str, float]

get_last_weights(n=1)

Return type:

dict[str, float]

format_loss(is_loss, n)

Return type:

dict[str, tuple[float, float]]

update_scheduler(schedulers, it)

Return type:

Scheduler

class konfai.network.network.ModuleArgsDict

Bases: Module, ABC

Named module graph container supporting KonfAI branch routing metadata.

class ModuleArgs(in_branch, out_branch, pretrained, alias, requires_grad, training)

Bases: object

in_channels: int | None

in_is_channel: bool

out_channels: int | None

out_is_channel: bool

keys()

Return type:

Iterable[str]

items()

Return type:

Iterable[tuple[str, Module | None]]

values()

Return type:

Iterable[Module | None]

add_module(name, module, in_branch=[0], out_branch=[0], pretrained=True, alias=[], requires_grad=None, training=None)

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

• name (str) – name of the child module. The child module can be accessed from this module using the given name

• module (Module) – child module to be added to the module.

Return type:

None

get_mapping()

static init_func(module, init_type, init_gain)

init(init_type, init_gain)

named_forward(*inputs, attributes=None)

Return type:

Iterator[tuple[str, Tensor]]

forward(*input)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

named_parameters(pretrained=False, recurse=False)

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

• prefix (str) – prefix to prepend to all parameter names.

• recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

• remove_duplicate (bool, optional) – whether to remove the duplicated parameters in the result. Defaults to True.

Yields:

(str, Parameter) – Tuple containing the name and parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

parameters(pretrained=False)

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

Yields:

Parameter – module parameter

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

named_module_args_dict()

Return type:

Iterator[tuple[str, Self, ModuleArgs]]

class konfai.network.network.OutputsGroup(measure)

Bases: list

Container describing one model output and its source modules.

layers: dict[str, Tensor]

add_layer(name, layer)

is_done()

clear()

Remove all items from list.

class konfai.network.network.Network(in_channels=1, optimizer=None, schedulers=None, outputs_criterions=None, patch=None, nb_batch_per_step=1, init_type='normal', init_gain=0.02, dim=3)

Bases: ModuleArgsDict, ABC

Base class for KonfAI networks participating in a routed model graph.

optimizer: Optimizer | None

schedulers: dict[_LRScheduler, int]

measure: Measure | None

outputsGroup: list[OutputsGroup]

state_dict(*args, **kwargs) → dict[str, object]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Note

The returned object is a shallow copy. It contains references to the module’s parameters and buffers.

Warning

Currently state_dict() also accepts positional arguments for destination, prefix and keep_vars in order. However, this is being deprecated and keyword arguments will be enforced in future releases.

Warning

Please avoid the use of argument destination as it is not designed for end-users.

Parameters:

• destination (dict, optional) – If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an OrderedDict will be created and returned. Default: None.

• prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''.

• keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching will not be performed. Default: False.

Returns:

a dictionary containing a whole state of the module

Return type:

dict[str, object]

Example:

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

load_state_dict(state_dict)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

Warning

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

Parameters:

• state_dict (dict[str, Tensor]) – a dict containing parameters and persistent buffers.

• strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Parameter for which the value from the module is preserved. Default: False

Returns:

• missing_keys is a list of str containing any keys that are expected

  by this module but missing from the provided state_dict.

• unexpected_keys is a list of str containing the keys that are not

  expected by this module but present in the provided state_dict.

Return type:

NamedTuple with missing_keys and unexpected_keys fields

Note

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a RuntimeError.

apply(fn)

Apply fn to each non-KonfAI child module and finally to self.

This overrides torch.nn.Module.apply so the recursive traversal can skip nested Network instances and keep KonfAI’s graph semantics intact.

Return type:

None

load(*args, **kwargs) → dict[str, object]

Return type:

dict[str, object]

init(*args, **kwargs) → dict[str, object]

Return type:

dict[str, object]

initialized()

named_forward(*inputs, attributes=None)

Return type:

Iterator[tuple[str, Tensor]]

get_layers(inputs, layers_name, attributes=None)

Return type:

Iterator[tuple[str, Tensor, PatchIndexed | None]]

init_outputs_group()

forward(batch_sample, output_layers=[])

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

list[tuple[str, Tensor]]

reset_loss(*args, **kwargs) → dict[str, object]

Return type:

dict[str, object]

backward(*args, **kwargs) → dict[str, object]

Return type:

dict[str, object]

update_lr(*args, **kwargs) → dict[str, object]

Return type:

dict[str, object]

get_networks(*args, **kwargs) → dict[str, object]

Return type:

dict[str, object]

static to(module, device, _counter=None)

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

• device (int) – the desired device of the parameters and buffers in this module

• dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

• tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

• memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

get_name()

Return type:

str

set_name(name)

Return type:

Self

set_state(state)

class konfai.network.network.MinimalModel(model, optimizer=<konfai.network.network.OptimizerLoader object>, schedulers={'default|StepLR': <konfai.network.network.LRSchedulersLoader object>}, outputs_criterions={'default': <konfai.network.network.TargetCriterionsLoader object>}, patch=None, dim=3, nb_batch_per_step=1, init_type='normal', init_gain=0.02)

Bases: Network

Small wrapper exposing a single network as a full KonfAI model graph.

The wrapped model arrives fully constructed — possibly carrying pretrained weights (a torchvision/MONAI/SMP class with weights=...). load therefore never re-initialises: load(init=True) at training start applies init_func over every descendant and would silently destroy those weights with init_type noise. Models built from scratch keep KonfAI’s init behaviour; checkpoint loading is unaffected.

load(state_dict, init=True, ema=False, override_lr=None)

class konfai.network.network.ModelLoader(classpath='default|segmentation.UNet.UNet')

Bases: object

Instantiate the root model graph declared in the active configuration.

get_model(train=True, konfai_args=None, konfai_without=['optimizer', 'schedulers', 'nb_batch_per_step', 'init_type', 'init_gain'])

Return type:

Network

class konfai.network.network.Model(model)

Bases: object

High-level model wrapper combining networks, criteria, and execution state.

train()

eval()

Module contents

Model graph primitives, blocks, and loaders used by KonfAI workflows.