noether.core.models.composite

Classes

CompositeModel

A composite model that consists of multiple submodels of type Model. By having multiple submodels, each model can have its own optimizer and learning rate scheduler, from weights etc.

Module Contents

class noether.core.models.composite.CompositeModel(model_config, update_counter=None, path_provider=None, data_container=None)

Bases: noether.core.models.base.ModelBase

A composite model that consists of multiple submodels of type Model. By having multiple submodels, each model can have its own optimizer and learning rate scheduler, from weights etc. This is useful for multi-component models,

A composite model must implement the submodels property, which returns a dictionary of submodel names to submodel instances.

Example code (dummy code):

from noether.core.models.composite import CompositeModel
from somewhere import MyModel1, MyModel2

class MyCompositeModel(CompositeModel):
    def __init__(self, model_config: MyCompositeModelConfig, update_counter: UpdateCounter | None = None, path_provider: PathProvider | None = None, data_container: DataContainer | None = None, static_context: dict[str, Any] | None = None):
        super().__init__(model_config, ...)

        self.submodel1 = MyModel1(
            model_config=model_config.submodel1_config,
            is_frozen=model_config.is_frozen,
            update_counter=update_counter,
            path_provider=path_provider,
            data_container=data_container,
            static_context=static_context,
            optimizer_config=model_config.submodel1_config.optimizer_config,
        )
        self.submodel2 = MyModel2(model_config=model_config.submodel2_config, ... )

    @property
    def submodels(self) -> dict[str, Model]:
        return dict(
            submodel1=self.submodel1,
            submodel2=self.submodel2,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # define forward pass here using self.submodel1 and self.submodel2
        x = self.submodel1(x)
        x = self.submodel2(x)
        return x

Base class for composite models, i.e. models that consist of multiple submodels of type Model.

Parameters:
property submodels: dict[str, noether.core.models.base.ModelBase]
Abstractmethod:

Return type:

dict[str, noether.core.models.base.ModelBase]

Returns the submodels of the composite model. This method must be implemented by the subclass, otherwise a NotImplementedError is raised.

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, noether.core.models.base.ModelBase]

property device: torch.device
Return type:

torch.device

initialize_weights()

Initialize the weights of the model, calling the initializer of all submodules.

Return type:

Self

apply_initializers()

Apply the initializers to the model, calling the initializer of all submodules.

Return type:

Self

initialize_optimizer()

Initialize the optimizer of the model.

Return type:

None

optimizer_step(grad_scaler)

Perform an optimization step, calling all submodules’ optimization steps.

Parameters:

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

Return type:

None

optimizer_schedule_step()

Perform the optimizer learning rate scheduler step, calling all submodules’ scheduler steps.

Return type:

None

optimizer_zero_grad(set_to_none=True)

Zero the gradients of the optimizer, calling all submodules’ zero_grad methods.

Parameters:

set_to_none (bool)

Return type:

None

property is_frozen: bool
Return type:

bool

train(mode=True)

Set the model to train or eval mode.

Overwrites the nn.Module.train method to avoid setting the model to train mode if it is frozen and to call all submodules’ train methods.

Parameters:

mode – If True, set the model to train mode. If False, set the model to eval mode.

Return type:

Self

to(device, *args, **kwargs)

Performs Tensor dtype and/or device conversion, calling all submodules’ to methods.

Parameters:

device – The desired device of the tensor. Can be a string (e.g. “cuda:0”) or “cpu”.

Return type:

Self