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.

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

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

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.Exit

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

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

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)

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)

Return type:

Iterator[tuple[str, Tensor]]

get_layers(inputs, layers_name)

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)

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.

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.