scaffold.data.artifact_manager

Submodules

Classes

ModelLogger

Class to be used for logging a model or state dicts as an artifact.

Package Contents

class scaffold.data.artifact_manager.ModelLogger(artifact_manager: scaffold.data.artifact_manager.base.ArtifactManager)

Class to be used for logging a model or state dicts as an artifact. Will always use a temporary dir locally to save artifacts before logging that dir.

Initializes the ModelLogger.

Parameters:

artifact_manager (ArtifactManager) – ArtifactManager to use for logging artifacts.

log_state_to_artifact(afid: str, model: torch.nn.Module, artifact_description: str, optimizers: List[torch.optim.Optimizer] = None, collection: str | None = None, **kwargs) → str

Logs the model and optimizer state dicts under the specified artifact id.

Parameters:

  • afid (str) – Artifact id to log under on the given artifact store or None to generate a new one.

  • model (torch.nn.Module) – Model to save the state dict off.

  • artifact_description (str) – Description of the artifact. Will be logged at <artifact_root>/ARTIFACT_META_DIR/ARTIFACT_DESCRIPTION_FILE and serves to reduce undocumented artifact clutter.

  • optimizers (List[torch.optim.Optimizer]) – All optimizers to save the state dicts off.

  • collection (Optional[str]) – Collection to log the artifact to.

  • kwargs – Additional key value pairs to save to the state dict

retrieve_state_from_artifact(afid: str, collection: str | None = None, version: int = None, device: str | None = None) → Dict

Returns a dictionary containing the model and optimizer state from the artifact store at the artifact_path.

Parameters:

  • afid – Artifact id of the model. Must have been logged via log_state_to_artifact()

  • device – Device to load the model to. “cpu”, “cuda” or None, which is the default that will use the gpu if available

Returns:

State dictionary containing the model and optimizer state. See log_state_to_artifact() for the key-value pairs.

static save_state(dict_dp: str, model: torch.nn.Module, optimizers: List[torch.optim.Optimizer] = None, schedulers: List[torch.optim.lr_scheduler._LRScheduler] = None, **kwargs) → None

Saves the model and optimizer state dicts, and optionally some keyword arguments, into one dictionary that can later be loaded to retrieve the exact state of the model training.

Notes

We intentionally do not save the model.state_dict() here since we need to support easy model loading from a single afid in the evaluator without initializing the model class first.

Parameters:

  • dict_dp (str) – Directory where the state dict pickle file should be saved.

  • model (torch.nn.Module) – Model to save the state dict off.

  • optimizers (List[torch.optim.Optimizer]) – All optimizers to save the state dicts off.

  • kwargs – Additional key value pairs to save to the state dict

artifact_manager