noether.core.models.base

Classes

ModelBaseConfig	Internal base class for all registry-based configs.
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.ModelBaseConfig(/, **data)

Bases: noether.core.schemas.lib._RegistryBase

Internal base class for all registry-based configs.

Provides auto-registration via __init_subclass__. Not meant to be used directly - use specific config base classes instead.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

kind: str | None = None

Kind of model to use, i.e. class path

name: str

Name of the model. Needs to be unique

optimizer_config: noether.core.optimizer.schemas.AnyOptimizerConfig | None = None

The optimizer configuration to use for training the model. When a model is used for inference only, this can be left as None.

initializers: list[Annotated[noether.core.initializers.AnyInitializer, Field(discriminator=kind)]] | None = None

List of initializers configs to use for the model.

is_frozen: bool | None = False

Whether to freeze the model parameters (i.e., not trainable).

forward_properties: list[str] | None = []

List of properties to be used as inputs for the forward pass of the model. Only relevant when the train_step of the BaseTrainer is used. When overridden in a class method, this property is ignored.

model_config

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

property config_kind: str

The fully qualified import path for the configuration class.

Return type:

str

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