Base models API

DeepVisionModel(*, backend, pre, post, loss_fn, device=None, pred_params=None, use_amp=None, caps=None, autocast_dtype=torch.float16, builder=None)

Bases: torch.nn.Module

High-level model façade used by deepvisiontools. All available wrappers are DeepVisionModels. This façade exposes public methods : train_step, eval_step, predict to run model in/without targets presence. This façade allows to save checkpoints with all needed parametrization.

Source code in src/deepvisiontools/models/base/basemodel.py

def __init__(
    self,
    *,
    backend: Backend,
    pre: Preprocessor,
    post: Postprocessor,
    loss_fn: BaseLoss,
    device: str | torch.device | None = None,
    pred_params: BasePredParams | None = None,
    use_amp: bool | None = None,
    caps: Capabilities | None = None,
    autocast_dtype: torch.dtype = torch.float16,
    builder: str | Callable | None = None,
) -> None:
    super().__init__()
    self.backend = backend
    self.pre = pre
    self.post = post
    self.loss_fn = loss_fn
    self.caps = caps or Capabilities(task="unknown")
    self.autocast_dtype = autocast_dtype
    self.pred_params = pred_params

    cfg_amp = bool(getattr(current(), "amp", True))
    self.use_amp: bool = cfg_amp if use_amp is None else bool(use_amp)

    # Move entire module tree (including backend/pre/post/loss_fn if they register params/buffers)
    self.to(_resolve_device(device))

train_step(images, targets=None, *, extra_args=None)

Generic train step: - targets are OPTIONAL - extra_args are forwarded to backend; if targets exist, we inject the preprocessed targets into extra_args["targets"] so facades that require targets during forward (Cascade Mask/RCNN) can use them.

Source code in src/deepvisiontools/models/base/basemodel.py

def train_step(
    self,
    images: torch.Tensor,
    targets: BatchData | None = None,
    *,
    extra_args: dict[str, Any] | None = None,
) -> ModelOutput:
    """
    Generic train step:
    - targets are OPTIONAL
    - extra_args are forwarded to backend; if targets exist, we inject the *preprocessed* targets
      into extra_args["targets"] so facades that require targets during forward (Cascade Mask/RCNN) can use them.
    """
    self.train()

    # ---- critical fix: move inputs to model device ----
    images, targets = self._move_images_targets(images, targets)

    ex: dict[str, Any] | None = (
        dict(extra_args) if isinstance(extra_args, dict) else None
    )

    if targets is None:
        xi = self.pre.images_only(images)
        with self._ac():
            raw = self.backend(xi, extra_args=ex)

        loss_dict = None
        try:
            loss_dict = self.loss_fn(raw, None)
        except Exception:
            loss_dict = None

        return ModelOutput(loss=loss_dict, preds=None, aux=raw)

    xi, yi = self.pre.images_targets(images, targets)

    if ex is None:
        ex = {}
    ex.setdefault("targets", yi)

    with self._ac():
        raw = self.backend(xi, extra_args=ex)

    loss_dict = self.loss_fn(raw, yi)
    return ModelOutput(loss=loss_dict, preds=None, aux=raw)

save(path, *, builder=None, model_kwargs=None, mode='deepvisiontools')

Save a portable model artifact for inference.

Parameters:

Name	Type	Description	Default
path	str	path to saved model. Extension must be .pt or .pth unless mode="safetensors".	required
builder	Optional[str]	if specific builder reconstruction is necessary. If None derives it automatically. Defaults to None.	None
model_kwargs	Optional[Dict[str, typing.Any]]	if specific builder kwargs are requested. If None derives it automatically. Defaults to None.	None
mode	Literal["deepvisiontools", "safetensors"]	if need to save backend state_dict only as safetensors (hugging face API compatible). Defaults to "deepvisiontools".	'deepvisiontools'

Note

- we recommend to use everything by default, unless you have something very specific.
- saftensors mode is convenient for hugging face deposit. A full json builder config is also provided to rebuild model.

Source code in src/deepvisiontools/models/base/basemodel.py

def save(
    self,
    path: str,
    *,
    builder: str | None = None,
    model_kwargs: dict[str, Any] | None = None,
    mode: Literal["deepvisiontools", "safetensors"] = "deepvisiontools",
) -> None:
    """Save a portable model artifact for inference.

    Args:
        path (str): path to saved model. Extension must be .pt or .pth unless `mode="safetensors"`.
        builder (Optional[str], optional): if specific builder reconstruction is necessary. If None derives it automatically. Defaults to None.
        model_kwargs (Optional[Dict[str, Any]], optional): if specific builder kwargs are requested. If None derives it automatically. Defaults to None.
        mode (Literal["deepvisiontools", "safetensors"], optional): if need to save backend state_dict only as safetensors (hugging face API compatible). Defaults to "deepvisiontools".

    ???+ note "Note"
        ```
        - we recommend to use everything by default, unless you have something very specific.
        - saftensors mode is convenient for hugging face deposit. A full json builder config is also provided to rebuild model.
        ```

    """

    if builder is None or model_kwargs is None:
        spec = self.export_spec()
        if spec is None:
            raise RuntimeError(
                "DeepVisionModel.save(): no export_spec available. "
                "Call model.set_export_spec(builder, model_kwargs) in the builder, "
                "or pass builder=... and model_kwargs=... explicitly."
            )
        if builder is None:
            builder = spec[0]
        if model_kwargs is None:
            model_kwargs = spec[1]

    if not isinstance(builder, str) or not builder:
        raise RuntimeError("DeepVisionModel.save(): invalid builder.")
    if not isinstance(model_kwargs, dict):
        raise RuntimeError("DeepVisionModel.save(): model_kwargs must be a dict.")

    normalized_mode = str(mode).strip().lower()

    if normalized_mode == "deepvisiontools":
        save_checkpoint(
            path,
            self,
            builder=str(builder),
            model_kwargs=dict(model_kwargs),
            epoch=None,
            optimizer=None,
            scaler=None,
            ema_state_dict=None,
            extra=None,
        )
        return

    if normalized_mode == "safetensors":
        save_safetensors_artifact(
            path,
            self,
            builder=str(builder),
            model_kwargs=dict(model_kwargs),
            extra=None,
            state_dict_target="backend",
        )
        return

    raise ValueError(
        f"Unsupported save mode '{mode}'. Expected 'deepvisiontools' or 'safetensors'."
    )

ModelFactory

Factory object used to build or restore a DeepVisionModel.

ModelFactory(...) is intentionally implemented through __new__: calling the class returns a ready-to-use model wrapper directly instead of returning a persistent factory instance.

Call parameters

ModelFactory(...) accepts the following call-time parameters:

• name: Registered model wrapper name, such as "smp", "yolo" or "yoloseg".
• cfg: Configuration dictionary with keys "name" and "kwargs".
• checkpoint: Path to a portable checkpoint saved with the deepvisiontools checkpoint helpers.
• weights: Tuple (builder_name, weights_path) used to build a model and load a state dictionary.
• **build_kwargs: Extra keyword arguments forwarded either to the selected builder or to the loading helpers, depending on the selected construction path.

For specific builder kwargs check documentation for the various model wrappers (build_yolo, build_smp, etc.).

Precedence

Calling ModelFactory(...) follows this precedence:

• if checkpoint is provided, restore the model from the saved builder/kwargs and checkpoint state;
• elif weights is provided, build through MODEL_REGISTRY and load the provided state dictionary;
• elif cfg is provided, build with cfg["name"] and cfg["kwargs"];
• else, name must be provided and remaining keyword arguments are forwarded to the selected builder.

All builders return a DeepVisionModel façade exposing train_step, eval_step, predict, preprocessing, postprocessing and loss handling.

Usage

from deepvisiontools.models import ModelFactory

model = ModelFactory(
    "smp",
    arch="Unet",
    encoder_name="resnet18",
    encoder_weights="imagenet",
    in_channels=3,
    num_classes=1,
)

cfg = {
    "name": "smp",
    "kwargs": {
        "arch": "Unet",
        "encoder_name": "resnet18",
        "encoder_weights": "imagenet",
        "in_channels": 3,
        "num_classes": 3,
    },
}
model = ModelFactory.from_config(cfg)

model = ModelFactory.from_checkpoint(
    "/path/to/checkpoint.pt",
    map_location="cuda",
    strict=False,
)

model = ModelFactory.from_weights(
    ("smp", "/path/to/weights.pth"),
    device="cuda",
    use_amp=True,
    strict=False,
    map_location="cpu",
    arch="Unet",
    encoder_name="resnet18",
    encoder_weights=None,
    in_channels=3,
    num_classes=1,
)

from_config(cfg) classmethod

Generate model from given config.

Parameters:

Name	Type	Description	Default
cfg	Dict[str, typing.Any]	config to be used.	required

Returns:

Type	Description
deepvisiontools.models.base.basemodel.DeepVisionModel	DeepVisionModel

Source code in src/deepvisiontools/models/base/basemodel.py

@classmethod
def from_config(cls, cfg: dict[str, Any]) -> DeepVisionModel:
    """Generate model from given config.

    Args:
        cfg (Dict[str, Any]): config to be used.

    Returns:
        DeepVisionModel
    """
    return cast(DeepVisionModel, cls(None, cfg=cfg))

from_checkpoint(path, *, map_location='cpu', strict=True) classmethod

Rebuild model from given checkpoint.

Parameters:

Name	Type	Description	Default
path	str	Path to checkpoint.	required
map_location	str, **optional**	device of model. Defaults to "cpu".	'cpu'
strict	bool, **optional**	requires strict state_dict matching. Defaults to True.	True

Returns:

Type	Description
deepvisiontools.models.base.basemodel.DeepVisionModel	DeepVisionModel

Source code in src/deepvisiontools/models/base/basemodel.py

@classmethod
def from_checkpoint(
    cls, path: str, *, map_location="cpu", strict=True
) -> DeepVisionModel:
    """Rebuild model from given checkpoint.

    Args:
        path (str): Path to checkpoint.
        map_location (str, **optional**): device of model. Defaults to "cpu".
        strict (bool, **optional**): requires strict state_dict matching. Defaults to True.

    Returns:
        DeepVisionModel
    """
    return cast(
        DeepVisionModel,
        cls(None, checkpoint=path, map_location=map_location, strict=strict),
    )

from_weights(builder_and_path, *, device=None, use_amp=None, strict=True, map_location=None, **build_kwargs) classmethod

Instantiate a model using the registry, then load a raw state_dict.

Supported raw weight formats: - torch files (.pt / .pth / etc.): interpreted as full-wrapper state_dict - safetensors (.safetensors): interpreted as backend-only state_dict

Notes: - This method is for loading raw weights when the caller already knows how the target model must be built (via builder + build_kwargs). - For portable DeepVision artifacts with sidecar metadata, prefer from_checkpoint(...), which can reconstruct the model from saved metadata.

Source code in src/deepvisiontools/models/base/basemodel.py

@classmethod
def from_weights(
    cls,
    builder_and_path,  # Tuple[str, str]
    *,
    device=None,
    use_amp=None,
    strict=True,
    map_location=None,
    **build_kwargs,
):
    """
    Instantiate a model using the registry, then load a raw state_dict.

    Supported raw weight formats:
    - torch files (.pt / .pth / etc.): interpreted as full-wrapper state_dict
    - safetensors (.safetensors): interpreted as backend-only state_dict

    Notes:
    - This method is for loading raw weights when the caller already knows how the
    target model must be built (via builder + build_kwargs).
    - For portable DeepVision artifacts with sidecar metadata, prefer
    `from_checkpoint(...)`, which can reconstruct the model from saved metadata.
    """
    from pathlib import Path

    build_kwargs = dict(build_kwargs)

    if "device" in build_kwargs:
        if device is None:
            device = build_kwargs.pop("device")
        else:
            build_kwargs.pop("device")

    if "use_amp" in build_kwargs:
        if use_amp is None:
            use_amp = build_kwargs.pop("use_amp")
        else:
            build_kwargs.pop("use_amp")

    bname, wpath = builder_and_path

    model = MODEL_REGISTRY.build(
        bname,
        device=device,
        use_amp=use_amp,
        **build_kwargs,
    )

    sd = load_raw_state_dict(
        wpath,
        map_location=("cpu" if map_location is None else map_location),
    )

    suffix = Path(wpath).suffix.lower()
    if suffix == ".safetensors":
        backend = getattr(model, "backend", None)
        if backend is None or not isinstance(backend, torch.nn.Module):
            raise RuntimeError(
                "ModelFactory.from_weights(): received a '.safetensors' file, "
                "which is interpreted as backend-only weights, but the built model "
                "does not expose a torch.nn.Module 'backend'."
            )
        if not strict:
            sd = _prune_incompatible_keys_for_lenient_load(backend, sd)
            backend.load_state_dict(sd, strict=False)
        else:
            backend.load_state_dict(sd, strict=True)
    else:
        if not strict:
            sd = _prune_incompatible_keys_for_lenient_load(model, sd)
            model.load_state_dict(sd, strict=False)
        else:
            model.load_state_dict(sd, strict=True)

    return model

available_configs() staticmethod

Return available config presets for each registered model wrapper.

The returned object is pretty-printable and can also be converted to a nested dict via .to_dict().

Output shape (dict form): { "bbox": {"yolo": [{"yolo11n": "..."}, ...], ...}, "instance_mask": {...}, "semantic_mask": {...}, }

Notes

• Builders may optionally expose __dvt_available_configs__: Dict[str(task), Dict[str(config_name), str(description)]]
• This method does NOT change how models are built; it only introspects registry metadata.

Source code in src/deepvisiontools/models/base/basemodel.py

@staticmethod
def available_configs() -> AvailableConfigs:
    """Return available config presets for each registered model wrapper.

    The returned object is pretty-printable and can also be converted to a
    nested dict via `.to_dict()`.

    Output shape (dict form):
      {
        "bbox": {"yolo": [{"yolo11n": "..."}, ...], ...},
        "instance_mask": {...},
        "semantic_mask": {...},
      }

    Notes:
      - Builders may optionally expose `__dvt_available_configs__`:
          Dict[str(task), Dict[str(config_name), str(description)]]
      - This method does NOT change how models are built; it only
        introspects registry metadata.
    """

    gathered: dict[str, dict[str, list[ConfigEntry]]] = {}

    builders = getattr(MODEL_REGISTRY, "_builders", {})
    if not isinstance(builders, dict):
        return AvailableConfigs(data={})

    for reg_name, builder in builders.items():
        meta = getattr(builder, "__dvt_available_configs__", None)
        if not isinstance(meta, dict):
            continue

        # Keep public registry key exactly as registered.
        wrapper_key = str(reg_name)

        for task, cfg_map in meta.items():
            if not isinstance(cfg_map, Mapping) or len(cfg_map) == 0:
                continue

            task_key = str(task)
            wrappers = gathered.setdefault(task_key, {})
            entries = wrappers.setdefault(wrapper_key, [])

            for cfg_name in sorted(cfg_map.keys()):
                desc = cfg_map.get(cfg_name, "")
                entries.append(
                    ConfigEntry(name=str(cfg_name), description=str(desc or ""))
                )

    for _, wrappers in gathered.items():
        for wrapper_key, entries in wrappers.items():
            wrappers[wrapper_key] = sorted(entries, key=lambda e: e.name)

    return AvailableConfigs(data=gathered)

load(path, *, map_location='cpu', strict=True) classmethod

Load a portable model artifact saved by DeepVisionModel.save(). User decides "raw vs ema" by choosing the file (e.g. model_best.pt vs model_best_ema.pt).

Source code in src/deepvisiontools/models/base/basemodel.py

@classmethod
def load(
    cls, path: str, *, map_location="cpu", strict: bool = True
) -> DeepVisionModel:
    """
    Load a portable model artifact saved by DeepVisionModel.save().
    User decides "raw vs ema" by choosing the file (e.g. model_best.pt vs model_best_ema.pt).
    """
    return cls.from_checkpoint(path, map_location=map_location, strict=strict)

load_checkpoint(path, builder, map_location='cpu', strict=True)

Load a model from a checkpoint saved by save_checkpoint, rebuild the architecture with the recorded builder name + kwargs, load weights (strict or lenient), and restore the model's extra state via set_extra_state when available.

The loader tolerates both the current flat format and an older nested meta format.

Parameters:

Name	Type	Description	Default
path	str	Path to the checkpoint produced by save_checkpoint.	required
builder	str | collections.abc.Callable	Either: - a callable of the form build(name, **kwargs) (e.g., MODEL_REGISTRY.build), or - a registry-like object that implements the same call signature. The function will be called using the stored (builder_name, model_kwargs).	required
map_location	str | torch.device	torch.load map location (e.g., "cpu", "cuda", torch.device("cuda:0")). Defaults to "cpu".	'cpu'
strict	bool	If True, enforce exact key matching for model.load_state_dict. If False, perform a lenient load by first calling _prune_incompatible_keys_for_lenient_load and then loading with strict=False. Defaults to True.	True

Returns:

Type	Description
tuple[torch.nn.Module, dict]	Tuple[torch.nn.Module, dict]: (model, meta) where: - model is rebuilt and weight-loaded, - meta contains at least: {"version": str, "builder": str, "model_kwargs": dict}.

Usage

# Using a registry builder callable
model, meta = load_checkpoint(
    "/tmp/ckpt.pt",
    MODEL_REGISTRY.build,
    map_location="cuda",
    strict=False,  # allow minor mismatches
)
print(meta["builder"], meta["model_kwargs"]["arch"])  # doctest: +ELLIPSIS
# After load, model extras (e.g., prediction params) are restored automatically
_ = model.eval()

Source code in src/deepvisiontools/models/base/checkpoints.py

def load_checkpoint(
    path: str,
    builder: str | Callable,
    map_location: str | torch.device = "cpu",
    strict: bool = True,
) -> tuple[torch.nn.Module, dict]:
    """
    Load a model from a checkpoint saved by `save_checkpoint`, **rebuild** the architecture
    with the recorded builder name + kwargs, load weights (strict or lenient), and restore
    the model's extra state via `set_extra_state` when available.

    The loader tolerates both the current flat format and an older nested `meta` format.

    Args:
        path (str): Path to the checkpoint produced by `save_checkpoint`.
        builder (str | Callable): Either:
            - a callable of the form `build(name, **kwargs)` (e.g., `MODEL_REGISTRY.build`), or
            - a registry-like object that implements the same call signature.
            The function will be called using the stored `(builder_name, model_kwargs)`.
        map_location (str | torch.device, optional): `torch.load` map location
            (e.g., `"cpu"`, `"cuda"`, `torch.device("cuda:0")`). Defaults to `"cpu"`.
        strict (bool, optional): If `True`, enforce exact key matching for
            `model.load_state_dict`. If `False`, perform a lenient load by first
            calling `_prune_incompatible_keys_for_lenient_load` and then loading
            with `strict=False`. Defaults to `True`.

    Returns:
        Tuple[torch.nn.Module, dict]: `(model, meta)` where:
            - `model` is rebuilt and weight-loaded,
            - `meta` contains at least: `{"version": str, "builder": str, "model_kwargs": dict}`.

    ???+ example "Usage"
        ```Python
        # Using a registry builder callable
        model, meta = load_checkpoint(
            "/tmp/ckpt.pt",
            MODEL_REGISTRY.build,
            map_location="cuda",
            strict=False,  # allow minor mismatches
        )
        print(meta["builder"], meta["model_kwargs"]["arch"])  # doctest: +ELLIPSIS
        # After load, model extras (e.g., prediction params) are restored automatically
        _ = model.eval()
        ```
    """
    if _is_safetensors_path(path):
        return load_safetensors_artifact(
            path,
            builder=builder,
            map_location=map_location,
            strict=strict,
        )

    pkg = torch.load(path, map_location=map_location)

    if "builder" in pkg and "model_kwargs" in pkg:
        model_name = pkg["builder"]
        model_kwargs = dict(pkg.get("model_kwargs", {}))
    else:
        meta = pkg.get("meta", {})
        b = meta.get("builder", {})
        model_name = b.get("name")
        model_kwargs = dict(b.get("kwargs", {}))

    if model_name is None:
        raise KeyError("Checkpoint missing 'builder' info (model name).")

    build_fn = cast(Callable[..., torch.nn.Module], builder)
    model = build_fn(model_name, **model_kwargs)

    state_dict = pkg.get("state_dict", None)
    if state_dict is None:
        state_dict = {k: v for k, v in pkg.items() if isinstance(v, torch.Tensor)}

    if not strict:
        pruned = _prune_incompatible_keys_for_lenient_load(model, state_dict)
        model.load_state_dict(pruned, strict=False)
    else:
        model.load_state_dict(state_dict, strict=True)

    extra = pkg.get("model_extra", None)
    if extra is not None and hasattr(model, "set_extra_state"):
        model.set_extra_state(extra)

    meta = {
        "version": pkg.get("version", pkg.get("_DEFAULT_VERSION", "1.0.0")),
        "builder": model_name,
        "model_kwargs": model_kwargs,
        "state_dict_target": "model",
    }
    return model, meta

load_training_state(path, registry_build, map_location='cpu', strict=True)

Load a training checkpoint saved by save_training_state / save_checkpoint and rebuild the model.

This function is a thin wrapper over load_checkpoint that returns a dict with the rebuilt model and a compact checkpoint meta.

Parameters:

Name	Type	Description	Default
path	str	Path to the saved training checkpoint.	required
registry_build	collections.abc.Callable	Callable of the form build(name, **kwargs) used to instantiate the model architecture (e.g., MODEL_REGISTRY.build).	required
map_location	str | torch.device	Map location for torch.load. Defaults to "cpu".	'cpu'
strict	bool	Whether to enforce strict key matching for model.load_state_dict. Defaults to True.	True

Returns:

Type	Description
dict[str, typing.Any]	Dict[str, Any]: {"model": torch.nn.Module, "ckpt": dict} where ckpt includes
dict[str, typing.Any]	at least "version", "builder", and "model_kwargs".

Usage

out = load_training_state("/tmp/train_ckpt.pt", MODEL_REGISTRY.build, map_location="cuda")
model = out["model"]
meta = out["ckpt"]
print(meta["builder"], meta["model_kwargs"]["arch"])  # doctest: +ELLIPSIS

Source code in src/deepvisiontools/models/base/checkpoints.py

def load_training_state(
    path: str,
    registry_build: Callable,
    map_location: str | torch.device = "cpu",
    strict: bool = True,
) -> dict[str, Any]:
    """
    Load a training checkpoint saved by `save_training_state` / `save_checkpoint` and
    rebuild the model.

    This function is a thin wrapper over `load_checkpoint` that returns a dict with
    the rebuilt model and a compact checkpoint meta.

    Args:
        path (str): Path to the saved training checkpoint.
        registry_build (Callable): Callable of the form `build(name, **kwargs)` used to
            instantiate the model architecture (e.g., `MODEL_REGISTRY.build`).
        map_location (str | torch.device, optional): Map location for `torch.load`.
            Defaults to `"cpu"`.
        strict (bool, optional): Whether to enforce strict key matching for
            `model.load_state_dict`. Defaults to `True`.

    Returns:
        Dict[str, Any]: `{"model": torch.nn.Module, "ckpt": dict}` where `ckpt` includes
        at least `"version"`, `"builder"`, and `"model_kwargs"`.

    ???+ example "Usage"
        ```Python
        out = load_training_state("/tmp/train_ckpt.pt", MODEL_REGISTRY.build, map_location="cuda")
        model = out["model"]
        meta = out["ckpt"]
        print(meta["builder"], meta["model_kwargs"]["arch"])  # doctest: +ELLIPSIS
        ```
    """
    model, ckpt = load_checkpoint(
        path, registry_build, map_location=map_location, strict=strict
    )
    return {"model": model, "ckpt": ckpt}

load_whole_object(path, map_location='cpu')

Load a whole-object pickle saved by save_whole_object.

⚠️ Prefer state-dict based checkpoints for robustness. This is intended for quick experiments where exact Python object restoration is acceptable.

Parameters:

Name	Type	Description	Default
path	str	Source file path pointing to a package produced by save_whole_object.	required
map_location	str | torch.device	Device mapping for torch.load (e.g., "cpu", "cuda"). Defaults to "cpu".	'cpu'

Returns:

Type	Description
torch.nn.Module	Tuple[torch.nn.Module, Dict[str, Any]]: (model, pkg) where pkg is the full
dict[str, typing.Any]	dictionary retrieved from disk (contains keys like "version", "object", "note").

Usage

obj, pkg = load_whole_object("/tmp/model_obj.pt", map_location="cpu")
isinstance(obj, torch.nn.Module)
>>> True
pkg.get("version")  # doctest: +ELLIPSIS

Source code in src/deepvisiontools/models/base/checkpoints.py

def load_whole_object(
    path: str, map_location: str | torch.device = "cpu"
) -> tuple[torch.nn.Module, dict[str, Any]]:
    """
    Load a whole-object pickle saved by `save_whole_object`.

    ⚠️ Prefer state-dict based checkpoints for robustness. This is intended for quick
    experiments where exact Python object restoration is acceptable.

    Args:
        path (str): Source file path pointing to a package produced by `save_whole_object`.
        map_location (str | torch.device, optional): Device mapping for `torch.load`
            (e.g., `"cpu"`, `"cuda"`). Defaults to `"cpu"`.

    Returns:
        Tuple[torch.nn.Module, Dict[str, Any]]: `(model, pkg)` where `pkg` is the full
        dictionary retrieved from disk (contains keys like `"version"`, `"object"`, `"note"`).

    ???+ example "Usage"
        ```Python
        obj, pkg = load_whole_object("/tmp/model_obj.pt", map_location="cpu")
        isinstance(obj, torch.nn.Module)
        >>> True
        pkg.get("version")  # doctest: +ELLIPSIS
        ```
    """
    pkg = torch.load(path, map_location=map_location)
    obj = pkg["object"]
    return obj, pkg

save_checkpoint(path, model, *, builder, model_kwargs, epoch=None, optimizer=None, scaler=None, ema_state_dict=None, extra=None, version=_DEFAULT_VERSION)

Save a training checkpoint suitable for later restoration with load_checkpoint.

Note

The checkpoint contains: - version string, - the builder name and its model_kwargs (to rebuild the model), - the model state_dict, - optional model_extra from model.get_extra_state() (e.g., prediction params), - training-time state: epoch, optimizer, AMP scaler, EMA weights, - arbitrary extra metadata, - RNG states (CPU and CUDA).

Parameters:

Name	Type	Description	Default
path	str	Destination file path (e.g., "/tmp/run123.pt").	required
model	torch.nn.Module	Model whose weights and extra state will be saved.	required
builder	str	Registry key or logical name used to rebuild the model later (e.g., "smp", "yolo", "yoloseg").	required
model_kwargs	Dict[str, typing.Any]	Exact kwargs that, together with builder, reproduce the model architecture when calling the registry builder.	required
epoch	Optional[int]	Current epoch index to store. Defaults to None.	None
optimizer	Optional[torch.optim.Optimizer]	Optimizer to serialize via optimizer.state_dict(). Defaults to None.	None
scaler	Optional[torch.cuda.amp.GradScaler]	AMP scaler; saved via scaler.state_dict(). Defaults to None.	None
ema_state_dict	Optional[Dict[str, typing.Any]]	Exponential Moving Average weights (already a state_dict). Defaults to None.	None
extra	Optional[Dict[str, typing.Any]]	Arbitrary metadata to persist (e.g., run config, metrics). Defaults to None.	None
version	str	Semantic or internal version tag for compatibility checks. Defaults to _DEFAULT_VERSION.	deepvisiontools.models.base.checkpoints._DEFAULT_VERSION

Returns:

Type	Description
None	None

Usage

save_checkpoint(
    "/tmp/ckpt.pt",
    model,
    builder="smp",
    model_kwargs=dict(
        arch="Unet",
        encoder_name="resnet18",
        in_channels=3,
        num_classes=1,
    ),
    epoch=12,
    optimizer=optim,
    scaler=scaler,      # optional
    ema_state_dict=ema.state_dict() if ema else None,
    extra={"val_iou": 0.71},
)

Source code in src/deepvisiontools/models/base/checkpoints.py

def save_checkpoint(
    path: str,
    model: torch.nn.Module,
    *,
    builder: str,
    model_kwargs: dict[str, Any],
    epoch: int | None = None,
    optimizer: torch.optim.Optimizer | None = None,
    scaler: torch.cuda.amp.GradScaler | None = None,
    ema_state_dict: dict[str, Any] | None = None,
    extra: dict[str, Any] | None = None,
    version: str = _DEFAULT_VERSION,
) -> None:
    """
    Save a **training checkpoint** suitable for later restoration with `load_checkpoint`.

    ???+ note "Note"
        The checkpoint contains:
        - version string,
        - the **builder name** and its `model_kwargs` (to rebuild the model),
        - the model `state_dict`,
        - optional `model_extra` from `model.get_extra_state()` (e.g., prediction params),
        - training-time state: `epoch`, `optimizer`, AMP `scaler`, EMA weights,
        - arbitrary `extra` metadata,
        - RNG states (CPU and CUDA).

    Args:
        path (str): Destination file path (e.g., `"/tmp/run123.pt"`).
        model (torch.nn.Module): Model whose weights and extra state will be saved.
        builder (str): Registry key or logical name used to rebuild the model later
            (e.g., `"smp"`, `"yolo"`, `"yoloseg"`).
        model_kwargs (Dict[str, Any]): Exact kwargs that, together with `builder`,
            reproduce the model architecture when calling the registry builder.
        epoch (Optional[int], optional): Current epoch index to store. Defaults to `None`.
        optimizer (Optional[torch.optim.Optimizer], optional): Optimizer to serialize via
            `optimizer.state_dict()`. Defaults to `None`.
        scaler (Optional[torch.cuda.amp.GradScaler], optional): AMP scaler; saved via
            `scaler.state_dict()`. Defaults to `None`.
        ema_state_dict (Optional[Dict[str, Any]], optional): Exponential Moving Average
            weights (already a `state_dict`). Defaults to `None`.
        extra (Optional[Dict[str, Any]], optional): Arbitrary metadata to persist
            (e.g., run config, metrics). Defaults to `None`.
        version (str, optional): Semantic or internal version tag for compatibility checks.
            Defaults to `_DEFAULT_VERSION`.

    Returns:
        None

    ???+ example "Usage"
        ```Python
        save_checkpoint(
            "/tmp/ckpt.pt",
            model,
            builder="smp",
            model_kwargs=dict(
                arch="Unet",
                encoder_name="resnet18",
                in_channels=3,
                num_classes=1,
            ),
            epoch=12,
            optimizer=optim,
            scaler=scaler,      # optional
            ema_state_dict=ema.state_dict() if ema else None,
            extra={"val_iou": 0.71},
        )
        ```
    """
    model_extra = None
    if hasattr(model, "get_extra_state"):
        try:
            model_extra = model.get_extra_state()
        except Exception:
            model_extra = None

    pkg: dict[str, Any] = {
        "version": version,
        "builder": builder,
        "model_kwargs": model_kwargs,
        "state_dict": model.state_dict(),
        "model_extra": model_extra,
        "epoch": epoch,
        "optimizer": optimizer.state_dict() if optimizer else None,
        "scaler": scaler.state_dict() if scaler else None,
        "ema": ema_state_dict,
        "extra": extra or {},
        "torch_rng": torch.get_rng_state(),
        "cuda_rng": (
            torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
        ),
    }
    torch.save(pkg, path)

save_training_state(path, model, optimizer, *, builder, model_kwargs, scheduler=None, scaler=None, ema=None, epoch=None, **meta)

Save a runnable training state that extends save_checkpoint with scheduler and EMA information. Intended for training restarts.

Parameters:

Name	Type	Description	Default
path	str	Destination file path.	required
model	torch.nn.Module	Model to save.	required
optimizer	torch.optim.Optimizer	Optimizer to serialize (state_dict()).	required
builder	str	Registry name used to rebuild the model (e.g., "smp").	required
model_kwargs	Dict[str, typing.Any]	Keyword arguments to rebuild the model.	required
scheduler	Optional[typing.Any]	LR scheduler; saved via state_dict() if provided. Defaults to None.	None
scaler	Optional[torch.cuda.amp.GradScaler]	AMP scaler. Defaults to None.	None
ema	Optional[torch.nn.Module]	EMA-wrapped model; its state_dict() is saved under the "ema" field. Defaults to None.	None
epoch	Optional[int]	Epoch index to persist. Defaults to None.	None
**meta	typing.Any	Arbitrary metadata merged into the "extra" field (e.g., metrics).	{}

Returns:

Type	Description
None	None

Usage

save_training_state(
    "/tmp/train_ckpt.pt",
    model,
    optimizer,
    builder="smp",
    model_kwargs={"arch": "Unet", "encoder_name": "resnet18", "num_classes": 1},
    scheduler=scheduler,
    scaler=scaler,
    ema=ema_model,
    epoch=5,
    best_val_iou=0.73,
)

Source code in src/deepvisiontools/models/base/checkpoints.py

def save_training_state(
    path: str,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    *,
    builder: str,
    model_kwargs: dict[str, Any],
    scheduler: Any | None = None,
    scaler: torch.cuda.amp.GradScaler | None = None,
    ema: torch.nn.Module | None = None,
    epoch: int | None = None,
    **meta: Any,
) -> None:
    """
    Save a **runnable training state** that extends `save_checkpoint` with scheduler
    and EMA information. Intended for training restarts.

    Args:
        path (str): Destination file path.
        model (torch.nn.Module): Model to save.
        optimizer (torch.optim.Optimizer): Optimizer to serialize (`state_dict()`).
        builder (str): Registry name used to rebuild the model (e.g., `"smp"`).
        model_kwargs (Dict[str, Any]): Keyword arguments to rebuild the model.
        scheduler (Optional[Any], optional): LR scheduler; saved via `state_dict()` if provided.
            Defaults to `None`.
        scaler (Optional[torch.cuda.amp.GradScaler], optional): AMP scaler. Defaults to `None`.
        ema (Optional[torch.nn.Module], optional): EMA-wrapped model; its `state_dict()` is saved
            under the `"ema"` field. Defaults to `None`.
        epoch (Optional[int], optional): Epoch index to persist. Defaults to `None`.
        **meta: Arbitrary metadata merged into the `"extra"` field (e.g., metrics).

    Returns:
        None

    ??? example "Usage"
        ```Python
        save_training_state(
            "/tmp/train_ckpt.pt",
            model,
            optimizer,
            builder="smp",
            model_kwargs={"arch": "Unet", "encoder_name": "resnet18", "num_classes": 1},
            scheduler=scheduler,
            scaler=scaler,
            ema=ema_model,
            epoch=5,
            best_val_iou=0.73,
        )
        ```
    """
    extra = {"scheduler": scheduler.state_dict() if scheduler else None, **meta}
    ema_sd = ema.state_dict() if ema else None
    save_checkpoint(
        path,
        model,
        builder=builder,
        model_kwargs=model_kwargs,
        epoch=epoch,
        optimizer=optimizer,
        scaler=scaler,
        ema_state_dict=ema_sd,
        extra=extra,
    )

save_whole_object(path, model, *, version=_DEFAULT_VERSION, note='')

Save the entire Python object (torch.nn.Module) via pickling.

⚠️ Not recommended for long-term portability—prefer save_checkpoint (weights + explicit rebuild info). Whole-object saves may break across library versions.

Parameters:

Name	Type	Description	Default
path	str	Destination file path.	required
model	torch.nn.Module	The model object to pickle.	required
version	str	Version tag written into the package. Defaults to _DEFAULT_VERSION.	deepvisiontools.models.base.checkpoints._DEFAULT_VERSION
note	str	Free-form note stored alongside the object. Defaults to "".	''

Returns:

Type	Description
None	None

Usage

save_whole_object("/tmp/model_obj.pt", model, note="dev snapshot")

Source code in src/deepvisiontools/models/base/checkpoints.py

def save_whole_object(
    path: str,
    model: torch.nn.Module,
    *,
    version: str = _DEFAULT_VERSION,
    note: str = "",
) -> None:
    """
    Save the **entire Python object** (`torch.nn.Module`) via pickling.

    ⚠️ Not recommended for long-term portability—prefer `save_checkpoint` (weights +
    explicit rebuild info). Whole-object saves may break across library versions.

    Args:
        path (str): Destination file path.
        model (torch.nn.Module): The model object to pickle.
        version (str, optional): Version tag written into the package. Defaults to
            `_DEFAULT_VERSION`.
        note (str, optional): Free-form note stored alongside the object. Defaults to `""`.

    Returns:
        None

    ???+ example "Usage"
        save_whole_object("/tmp/model_obj.pt", model, note="dev snapshot")
    """
    torch.save({"version": version, "object": model, "note": note}, path)