ksuit.models.model_base

Classes

ModelBase

Base class for all neural network modules.

Module Contents

class ksuit.models.model_base.ModelBase(model_config, update_counter=None, path_provider=None, data_container=None, initializer_config=None, static_context=None)

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

Base class for ksuit models that is used to define the interface for all models trainable by ksuit trainers.

Provides methods to initialize the model weights and setup (model-specific) optimizers.

Parameters:

  • model_config (ksuit.schemas.models.ModelBaseConfig) – Model configuration.

  • update_counter (ksuit.utils.training.counter.UpdateCounter | None) – The update counter provided to the optimizer.

  • path_provider (ksuit.providers.PathProvider | None) – A path provider used by the initializer to store or retrieve checkpoints.

  • data_container (ksuit.utils.data.data_container.DataContainer | None) – The data container which includes the data and dataloader. This is currently unused but helpful for quick prototyping only, evaluating forward in debug mode, etc.

  • initializer_config (ksuit.schemas.initializers.InitializerConfig | None) – The initializer config used to initialize the model e.g. from a checkpoint.

  • static_context (dict[str, Any] | None) – The static context used to pass information between submodules, e.g. patch_size, latent_dim.

logger
name
update_counter = None
path_provider = None
data_container = None
initializers: list[ksuit.initializers.InitializerBase] = []
static_context
model_config
is_initialized = False
property optimizer: ksuit.optimizer.OptimizerWrapper
Return type:

ksuit.optimizer.OptimizerWrapper

property device: torch.device
Abstractmethod:

Return type:

torch.device

property is_frozen: bool
Abstractmethod:

Return type:

bool

property param_count: int
Return type:

int

property trainable_param_count: int
Return type:

int

property frozen_param_count: int
Return type:

int

property nograd_paramnames: list[str]
Return type:

list[str]

initialize()

Initializes weights and optimizer parameters of the model.

abstractmethod get_named_models()

Returns a dict of {model_name: model}, e.g., to log all learning rates of all models/submodels.

Return type:

dict[str, ModelBase]

abstractmethod initialize_weights()

Initialize the weights of the model.

Return type:

Self

abstractmethod apply_initializers()

Apply the initializers to the model.

Return type:

Self

abstractmethod initialize_optimizer()

Initialize the optimizer of the model.

Return type:

None

abstractmethod optimizer_step(grad_scaler)

Perform an optimization step.

Parameters:

grad_scaler (torch.amp.grad_scaler.GradScaler | None)

Return type:

None

abstractmethod optimizer_schedule_step()

Perform the optimizer learning rate scheduler step.

Return type:

None

abstractmethod optimizer_zero_grad(set_to_none=True)

Zero the gradients of the optimizer.

Parameters:

set_to_none (bool)

Return type:

None