Skip to content

io

IO module.

Functions:

load_model_state 

load_model_state(
    save_epochs_dir: str | Path,
    epoch_num: int | None = None,
) -> Mapping

Load a model from disk.

Parameters:

  • save_epochs_dir (str | Path) –

    The director where the models are saved

  • epoch_num (int | None, default: None ) –

    The epoch number or if None/-1 the best state will be loaded, by default None

Returns:

  • Mapping –

    The loaded model state and possibly an optimizer state.

Raises:

  • ValueError –

    When the directory does not exist or the requested model is not available.

Source code in src/autoden/models/io.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def load_model_state(save_epochs_dir: str | Path, epoch_num: int | None = None) -> Mapping:
    """Load a model from disk.

    Parameters
    ----------
    save_epochs_dir : str | Path
        The director where the models are saved
    epoch_num : int | None, optional
        The epoch number or if None/-1 the best state will be loaded, by default None

    Returns
    -------
    Mapping
        The loaded model state and possibly an optimizer state.

    Raises
    ------
    ValueError
        When the directory does not exist or the requested model is not available.
    """
    epochs_base_path = Path(save_epochs_dir) / "weights"
    if not epochs_base_path.exists():
        raise ValueError(f"Directory of the model state {epochs_base_path} does not exist!")

    if epoch_num is None or epoch_num == -1:
        state_path = epochs_base_path / "weights.pt"
    else:
        state_path = epochs_base_path / f"weights_epoch_{epoch_num}.pt"
    if not state_path.exists():
        raise ValueError(f"Model state {state_path} does not exist!")

    print(f"Loading state path: {state_path}")
    return pt.load(state_path)

save_model 

save_model(
    dst_file: str | Path,
    model: Module,
    optim_state: Mapping | None = None,
    epoch_num: int = 0,
) -> None

Save a model and optionally its optimizer state to a file.

Parameters:

  • dst_file (str or Path) –

    The destination file path where the model will be saved.

  • model (Module) –

    The model to be saved. It must implement the SerializableModel protocol.

  • optim_state (Mapping, default: None ) –

    The state of the optimizer to be saved. Default is None.

  • epoch_num (int, default: 0 ) –

    The current epoch number. Default is 0.

Raises:

  • ValueError –

    If the model does not implement the SerializableModel protocol.

Source code in src/autoden/models/io.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def save_model(dst_file: str | Path, model: Module, optim_state: Mapping | None = None, epoch_num: int = 0) -> None:
    """
    Save a model and optionally its optimizer state to a file.

    Parameters
    ----------
    dst_file : str or Path
        The destination file path where the model will be saved.
    model : Module
        The model to be saved. It must implement the `SerializableModel` protocol.
    optim_state : Mapping, optional
        The state of the optimizer to be saved. Default is None.
    epoch_num : int, optional
        The current epoch number. Default is 0.

    Raises
    ------
    ValueError
        If the model does not implement the `SerializableModel` protocol.
    """
    if not isinstance(model, SerializableModel):
        raise ValueError("The model needs to implement the protocol SerializableModel, in order to be writable to disk")

    pt.save(
        {
            "model_class": model.__class__.__name__,
            "init_params": model.init_params,
            "epoch": epoch_num,
            "state_dict": model.state_dict(),
            "optimizer": optim_state,
        },
        dst_file,
    )

save_model_state 

save_model_state(
    save_epochs_dir: str | Path,
    epoch_num: int,
    model: Module,
    optim_state: Mapping | None = None,
    is_best: bool = False,
) -> None

Save a model's state to disk.

This function saves the state of a model and optionally its optimizer to disk. The model state is saved in a directory specified by save_epochs_dir. If is_best is True, the model state is saved as "weights.pt". Otherwise, it is saved with a filename that includes the epoch number.

Parameters:

  • save_epochs_dir (str | Path) –

    The directory where to save the model state.

  • epoch_num (int) –

    The epoch number.

  • model (Module) –

    The model whose state is to be saved.

  • optim_state (Mapping, default: None ) –

    The optimizer state to save, by default None.

  • is_best (bool, default: False ) –

    Whether it is the best fitted model, by default False.

Returns:

  • None –
Source code in src/autoden/models/io.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def save_model_state(
    save_epochs_dir: str | Path,
    epoch_num: int,
    model: Module,
    optim_state: Mapping | None = None,
    is_best: bool = False,
) -> None:
    """Save a model's state to disk.

    This function saves the state of a model and optionally its optimizer to disk.
    The model state is saved in a directory specified by `save_epochs_dir`. If
    `is_best` is True, the model state is saved as "weights.pt". Otherwise, it is
    saved with a filename that includes the epoch number.

    Parameters
    ----------
    save_epochs_dir : str | Path
        The directory where to save the model state.
    epoch_num : int
        The epoch number.
    model : Module
        The model whose state is to be saved.
    optim_state : Mapping, optional
        The optimizer state to save, by default None.
    is_best : bool, optional
        Whether it is the best fitted model, by default False.

    Returns
    -------
    None
    """
    epochs_base_path = Path(save_epochs_dir) / "weights"
    epochs_base_path.mkdir(parents=True, exist_ok=True)

    dst_file = epochs_base_path / ("weights.pt" if is_best else f"weights_epoch_{epoch_num}.pt")
    save_model(dst_file=dst_file, model=model, optim_state=optim_state, epoch_num=epoch_num)