Models API

KonfAI model graphs are configured through loaders, routed network containers, criteria, and reusable blocks.

Model graph and loaders

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

Bases: object

Instantiate the root model graph declared in the active configuration.

class konfai.network.network.Model(model)

Bases: object

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

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.

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

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

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)

class konfai.network.network.ModuleArgsDict

Bases: Module, ABC

Named module graph container supporting KonfAI branch routing metadata.

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

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)

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

Bases: object

Configuration-aware factory for PyTorch optimizers.

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

Bases: object

Configuration-aware factory for learning-rate schedulers.

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.

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.

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

Bases: object

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

Building blocks

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.

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.get_torch_module(name_fonction, dim=None)

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

Return type:

Module

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

Instantiate the normalization layer matching the requested mode.

Return type:

Module | None

Metric schedulers

class konfai.metric.schedulers.Scheduler(start_value)

Bases: object

Base class for scalar schedulers used by KonfAI criteria.

class konfai.metric.schedulers.Constant(value=1)

Bases: Scheduler

Scheduler returning a constant value for all iterations.

class konfai.metric.schedulers.CosineAnnealing(start_value=1, eta_min=1e-05, t_max=100)

Bases: Scheduler

Cosine annealing scheduler for criterion weights.

class konfai.metric.schedulers.Warmup(optimizer, warmup_steps=10, last_epoch=-1, verbose='deprecated')

Bases: LambdaLR

Learning-rate warmup wrapper compatible with PyTorch optimizers.

See also

• Model graph and output naming

• Using custom models, transforms, augmentations, and losses