Base Model Classes

Base classes for all models.

class openadmet.models.architecture.model_base.LightningModelBase[source]

Bases: ModelBase

A model that uses PyTorch Lightning.

freeze_weights(*args, **kwargs)[source]

Freeze parts of the model for transfer learning or fine-tuning.

Parameters:

  • *args (variable length argument list) – Arguments to be passed to the implementing model’s freeze_weights method.

  • **kwargs (keyword arguments) – Keyword arguments to be passed to the implementing model’s freeze_weights method.

Notes

This method should set the requires_grad attribute of the specified layers to False, preventing their weights from being updated during training. It also should set these layers to evaluation mode.

load(path: PathLike)[source]

Load the model from a file.

Parameters:

path (PathLike) – Path to load the model from

make_new()[source]

Copy parameters to a new model instance without copying the estimator.

Returns:

A new instance of LightningModelBase with the same parameters.

Return type:

LightningModelBase

model_config: ClassVar[ConfigDict] = {}

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

model_post_init(context: Any, /) → None

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Args:

self: The BaseModel instance. context: The context.

save(path: PathLike)[source]

Save the model to a file.

Parameters:

path (PathLike) – Path to save the model to

serialize(param_path: PathLike = 'model.json', serial_path: PathLike = 'model.pth')[source]

Save the model to a json file and a serialized file.

Parameters:

  • param_path (PathLike) – Path to save the model parameters to

  • serial_path (PathLike) – Path to save the serialized model to

type: ClassVar[str]
class openadmet.models.architecture.model_base.LightningModuleBase(*args: Any, **kwargs: Any)[source]

Bases: LightningModule

Lightning module base class.

A PyTorch lightning model may inherit this instead of pl.LightningModule to preconfigure optimizer and scheduler.

configure_optimizers()[source]

Return optimizer and scheduler configuration for Lightning’s configure_optimizers.

monitor_metric: str = 'val_loss'
optimizer: str = 'adamw'
optimizer_lr: float = 0.001
optimizer_weight_decay: float = 1e-05
scheduler: str = 'cosine'
scheduler_factor: float = 0.5
scheduler_patience: int = 10
type: ClassVar[str]
class openadmet.models.architecture.model_base.ModelBase[source]

Bases: BaseModel, ABC

Base class for all models.

abstract build()[source]

Prepare the model, abstract method to be implemented by subclasses.

abstract deserialize(param_path: PathLike, serial_path: PathLike)[source]

Deserialize the model, abstract method to be implemented by subclasses.

Parameters:

  • param_path (PathLike) – Path to load the model parameters from

  • serial_path (PathLike) – Path to load the model serialization from

property estimator

Get the model estimator.

abstract load(path: PathLike)[source]

Load the model, abstract method to be implemented by subclasses.

Parameters:

path (PathLike) – Path to load the model from

model_config: ClassVar[ConfigDict] = {}

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

model_post_init(context: Any, /) → None

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Args:

self: The BaseModel instance. context: The context.

abstract predict(input: Any)[source]

Predict using the model, abstract method to be implemented by subclasses.

Parameters:

input (Any) – Input data to predict on

abstract save(path: PathLike)[source]

Save the model, abstract method to be implemented by subclasses.

Parameters:

path (PathLike) – Path to save the model to

abstract serialize(param_path: PathLike, serial_path: PathLike)[source]

Serialize the model, abstract method to be implemented by subclasses.

Parameters:

  • param_path (PathLike) – Path to save the model parameters to

  • serial_path (PathLike) – Path to save the model serialization to

abstract train()[source]

Train the model, abstract method to be implemented by subclasses.

class openadmet.models.architecture.model_base.PickleableModelBase[source]

Bases: ModelBase

An sklearn model that can be pickled using joblib.

load(path: PathLike)[source]

Load the model from a pickle file.

Parameters:

path (PathLike) – Path to load the model from

make_new() → PickleableModelBase[source]

Copy parameters to a new model instance without copying the estimator.

model_config: ClassVar[ConfigDict] = {}

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

model_post_init(context: Any, /) → None

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Args:

self: The BaseModel instance. context: The context.

pickleable: ClassVar[bool] = True
save(path: PathLike)[source]

Save the model to a pickle file.

Parameters:

path (PathLike) – Path to save the model to

serialize(param_path: PathLike = 'model.json', serial_path: PathLike = 'model.pkl')[source]

Save the model to a json file and a pickled file.

Parameters:

  • param_path (PathLike) – Path to save the model parameters to

  • serial_path (PathLike) – Path to save the pickled model to

openadmet.models.architecture.model_base.get_mod_class(model_type)[source]

Get the model class from the registry.