Class BaseModel

climart/models/base_model.py:23–349

This is a template class, that should be inherited by any neural net emulator model. Methods that need to be implemented by your concrete NN model (just as if you would define a torch.nn.Module): - __init__(.) - forward(.) The other methods may be overridden as needed.

class BaseModel(LightningModule):
    """
    This is a template class, that should be inherited by any neural net emulator model.
    Methods that need to be implemented by your concrete NN model (just as if you would define a torch.nn.Module):
        - __init__(.)
        - forward(.)

    The other methods may be overridden as needed.
    It is recommended to define the attribute
        - self.example_input_array = torch.randn(<YourModelInputShape>)  # batch dimension can be anything, e.g. 7

    ------------
    A LightningModule organizes your PyTorch code into 5 sections:
        - Computations (init).
        - Train loop (training_step)
        - Validation loop (validation_step)
        - Test loop (test_step)
        - Optimizers (configure_optimizers)

    Read the docs:
        https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html
    """

    def __init__(self,
                 datamodule_config: DictConfig = None,
                 optimizer: Optional[DictConfig] = None,
                 scheduler: Optional[DictConfig] = None,
                 monitor: Optional[str] = None,
                 mode: str = "min",
                 loss_function: str = "mean_squared_error",
                 downwelling_loss_contribution: float = 0.5,
                 upwelling_loss_contribution: float = 0.5,
                 heating_rate_loss_contribution: float = 0.0,
                 input_transform: Optional[AbstractTransform] = None,
                 output_normalizer: Optional[Dict[str, NormalizationMethod]] = None,
                 out_layer_bias_init: Optional[np.ndarray] = None,
                 name: str = "",
                 verbose: bool = True,
                 ):
        super().__init__()
        self.log_text = get_logger(name=self.__class__.__name__ if name == '' else name)
        self.name = name
        self.verbose = verbose
        if not self.verbose:
            self.log_text.setLevel(logging.WARN)
        if input_transform is None or isinstance(input_transform, AbstractTransform):
            self.input_transform = input_transform
        else:
            self.input_transform = hydra.utils.instantiate(input_transform)
        if datamodule_config is not None:
            input_output_dimensions = get_data_dims(exp_type=datamodule_config.get("exp_type"))
            self.raw_input_dim = input_output_dimensions['input_dim']
            self.raw_output_dim = input_output_dimensions['output_dim']
            self.raw_spatial_dim = input_output_dimensions['spatial_dim']
            self.num_layers = self.raw_spatial_dim[LAYERS]
            self.num_levels = self.raw_spatial_dim[LEVELS]

        self.output_normalizer = output_normalizer.copy() if output_normalizer is not None else None