noether.core.models.model¶

Classes¶

Model

Model class that should be extended by all custom models.

Module Contents¶

class noether.core.models.model.Model(model_config, is_frozen=False, update_counter=None, path_provider=None, data_container=None)¶

Bases: noether.core.models.base.ModelBase

Model class that should be extended by all custom models. Each model has its own optimizer and learning rate scheduler, which are initialized in the initialize_optimizer method.

Example code (dummy code):

from noether.core.models.single import Model
from noether.core.schemas.models import ModelBaseConfig

class MyModelConfig(ModelBaseConfig):
    kind: path.to.MyModel
    name: my_model
    optimizer_config:
        kind: torch.optim.AdamW
        lr: 1.0e-3
        weight_decay: 0.05
        clip_grad_norm: 1.0
        schedule_config:
            kind: noether.core.schedules.LinearWarmupCosineDecaySchedule
            warmup_percent: 0.05
            end_value: 1.0e-6
            max_value: ${model.optimizer_config.lr}

    input_dim: int = 128
    hidden_dim: int = 256
    output_dim: int = 10

class MyModel(Model):
    def __init__(self, model_config: MyModelConfig, ...):
        super().__init__(model_config, ...)

        self.layer1 = torch.nn.Linear(self.model_config.input_dim, self.model_config.hidden_dim)
        self.layer2 = torch.nn.Linear(self.model_config.hidden_dim, self.model_config.output_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # define forward pass here
        x = self.layer1(x)
        x = torch.relu(x)
        x = self.layer2(x)
        return x

Base class for single models, i.e. one model with one optimizer as opposed to CompositeModel.

Parameters:

model_config (noether.core.schemas.models.ModelBaseConfig) – Model configuration. See ModelBaseConfig for available options.
update_counter (noether.core.utils.training.UpdateCounter | None) – The UpdateCounter provided to the optimizer.
is_frozen (bool) – If true, will set requires_grad of all parameters to false. Will also put the model into eval mode (e.g., to put Dropout or BatchNorm into eval mode).
path_provider (noether.core.providers.PathProvider | None) – PathProvider used by the initializer to store or retrieve checkpoints.
data_container (noether.data.container.DataContainer | None) – DataContainer which includes the data and dataloader. This is currently unused but helpful for quick prototyping only, evaluating forward in debug mode, etc.

property is_frozen: bool¶

Return type:: bool

property device: torch.device¶

Return type:: torch.device

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]

initialize_weights()¶

Freezes the weights of the model by setting requires_grad to False if self.is_frozen is True.

Return type:: Self

apply_initializers()¶

Apply the initializers to the model, calling initializer.init_weights and initializer.init_optim.

Return type:: Self

initialize_optimizer()¶

Initialize the optimizer.

Return type:: None

optimizer_step(grad_scaler)¶

Perform an optimization step.

Parameters:: grad_scaler (torch.amp.grad_scaler.GradScaler | None)
Return type:: None

optimizer_schedule_step()¶

Perform the optimizer learning rate scheduler step.

Return type:: None

optimizer_zero_grad(set_to_none=True)¶

Zero the gradients of the optimizer.

Parameters:: set_to_none (bool)
Return type:: None

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.

Parameters:: mode (bool) – 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, overwriting nn.Module.to method to set the _device attribute.

Parameters:: device (str | torch.device | int | None) – The desired device of the tensor. Can be a string (e.g. “cuda:0”) or “cpu”.
Return type:: Self