noether.core.models.base¶

Classes¶

ModelBase

Base class for models (Model and CompositeModel) that is used to define the interface for all models trainable by the BaseTrainer.

Module Contents¶

class noether.core.models.base.ModelBase(model_config, update_counter=None, path_provider=None, data_container=None, initializer_config=None)¶

Bases: torch.nn.Module

Base class for models (Model and CompositeModel) that is used to define the interface for all models trainable by the BaseTrainer.

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

Parameters:

model_config (noether.core.schemas.models.ModelBaseConfig) – Model configuration. See ModelBaseConfig for available options.
update_counter (noether.core.utils.training.counter.UpdateCounter | None) – The UpdateCounter provided to the optimizer.
path_provider (noether.core.providers.PathProvider | None) – A path PathProvider used by the initializer to store or retrieve checkpoints.
data_container (noether.data.container.DataContainer | None) – The DataContainer which includes the data and dataloader. This is currently unused but helpful for quick prototyping only, evaluating forward in debug mode, etc.
initializer_config (list[noether.core.schemas.initializers.InitializerConfig] | None) – The initializer config used to initialize the model e.g. from a checkpoint.

logger¶

name¶

update_counter = None¶

path_provider = None¶

data_container = None¶

initializers: list[noether.core.initializers.InitializerBase] = []¶

model_config¶

is_initialized = False¶

property optimizer: noether.core.optimizer.OptimizerWrapper | None¶

Return type:: noether.core.optimizer.OptimizerWrapper | None

property device: torch.device¶

Abstractmethod:
Return type:: torch.device

property is_frozen: bool¶

Abstractmethod:
Return type:: bool

property param_count: int¶

Returns the total number of parameters in the model.

Return type:: int

property trainable_param_count: int¶

Returns the number of parameters that require gradients (i.e., are trainable).

Return type:: int

property frozen_param_count: int¶

Returns the number of parameters that do not require gradients (i.e., are frozen).

Return type:: int

property nograd_paramnames: list[str]¶

Returns a list of parameter names that do not have gradients (i.e., grad is None) but require gradients.

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