Training API

Trainer(cfg)

Main class for training models in deepvisiontools. Everything configurable can be done through the TrainerConfig class.

Parameters:

Name	Type	Description	Default
cfg	deepvisiontools.train.trainconfig.TrainerConfig	Trainer configuration.	required

Usage

from deepvisiontools.config import current, set_current
from deepvisiontools.models import ModelFactory
from deepvisiontools.datasets import DeepVisionDataset, DeepVisionLoader
from deepvisiontools.train import Trainer, launch, TrainerConfig
from deepvisiontools.metrics.factories import (
    DetectionSuiteFactory,
    InstanceSegSuiteFactory,
    SemanticSegSuiteFactory,
)
from torch.optim import Adam
from torchvision.transforms.v2 import (
    RandomHorizontalFlip,
    RandomVerticalFlip,
    RandomGrayscale,
    ColorJitter,
)

# For a 1 class detection dataset, otherwise need to modify current() as explained in DVConfig class

# ── Datsets Paths ───────────────────────
TRAIN_DIR = "datasets_demo/instance_nc=1/train"  # e.g. /data/coco/train2017
VAL_DIR = "datasets_demo/instance_nc=1/valid"  # e.g. /data/coco/val2017

set_current(current().update_device("cuda"))

augs = [
    RandomVerticalFlip(),
    RandomVerticalFlip(),
    ColorJitter(hue=0.15, brightness=0.15, contrast=0.15, saturation=0.15),
    RandomGrayscale(0.05),
]

trainset = DeepVisionDataset(
    TRAIN_DIR,
    "coco",
    data_type_reader="instance_mask",
    augmentation=augs,
)
valset = DeepVisionDataset(
    VAL_DIR, "coco", data_type_reader="instance_mask"
)

trainloader = DeepVisionLoader(trainset, batch_size=4)
validloader = DeepVisionLoader(valset, batch_size=4)

# Best way to use deepvisiontools : build config from registered ones, then modify them if needed
Config = TrainerConfig.build_registered_config(
    model="detector-seg",
    config="detector_convnextv2-base_cascade_maskrcnn",
    train_loader=trainloader,
    valid_loader=validloader,
)
# exemple of modification but everything is configurable (optimizer, scheduler, models directly from ModelFactory etc.)
# Note : you can build TrainerConfig from scratch as well, but there are a bunch of parameters you may want to default to
Config.epochs = 120

# Start Training
def main():
    Trainer(Config).fit(resume=None)

launch(main, devices="auto")    # Can Use DDP wrapper if multiple gpu detected

# If you want to restart from a checkpoint :
Trainer.fit(resume="my/checkpoint/path") # in main

Source code in src/deepvisiontools/train/trainer.py

def __init__(self, cfg: TrainerConfig):
    cfg.validate()
    self.cfg = cfg

    # ----------------------------------
    # Device priority
    # ----------------------------------
    # Trainer cfg wins when explicitly set. If cfg.devices.device is None,
    # fall back to current().device at runtime.
    requested_device = _resolve_requested_device(
        getattr(cfg.devices, "device", None)
    )

    # ----------------------------------
    # Strategy priority
    # ----------------------------------
    # Explicit cfg.devices.strategy wins. "auto" means:
    #   - DDP only when launched under torchrun / launch(... nproc > 1)
    #   - otherwise single-device, even if multiple GPUs are visible
    strategy_name = _resolve_strategy_name(cfg.devices.strategy, requested_device)

    self.strategy: _Strategy
    if strategy_name == "ddp":
        self.strategy = DDPStrategy(
            requested_device=requested_device,
            find_unused_parameters=cfg.find_unused_parameters,
            timeout_seconds=getattr(cfg.devices, "ddp_timeout_seconds", None),
        )
    else:
        self.strategy = SingleDeviceStrategy()

    self.device = (
        self.strategy.resolve_device(requested_device)
        if hasattr(self.strategy, "resolve_device")
        else requested_device
    )

    # ----------------------------
    # Model (+ normalization policy BEFORE DDP wrap)
    # ----------------------------
    raw_model = cfg.build_model()
    raw_model, norm_summary = self._prepare_model_norm_policy(raw_model)

    self.model = self.strategy.setup(raw_model, self.device)
    self._norm_policy_summary: dict[str, Any] = dict(norm_summary)
    self._norm_policy_summary_emitted: bool = False

    m = cast(Any, _unwrap(self.model))
    try:
        if (
            self.cfg.checkpoint.model_builder is None
            or self.cfg.checkpoint.model_kwargs is None
        ) and hasattr(m, "export_spec"):
            spec = m.export_spec()
            if isinstance(spec, tuple) and len(spec) == 2:
                builder, model_kwargs = spec
                if self.cfg.checkpoint.model_builder is None:
                    self.cfg.checkpoint.model_builder = builder
                if self.cfg.checkpoint.model_kwargs is None:
                    self.cfg.checkpoint.model_kwargs = dict(model_kwargs)
    except Exception:
        pass

    # ----------------------------
    # Data
    # ----------------------------
    self.train_loader = self.strategy.wrap_loader(cfg.train_loader, shuffle=True)
    self.valid_loader = (
        self.strategy.wrap_loader(cfg.valid_loader, shuffle=False)
        if cfg.valid_loader is not None
        else None
    )
    self.train_loader = self._apply_loader_overrides(self.train_loader)
    self.valid_loader = self._apply_loader_overrides(self.valid_loader)

    # ----------------------------
    # Optimizer & Scheduler
    # ----------------------------
    param_groups = build_param_groups(_unwrap(self.model), cfg)
    self.optimizer = self._instantiate_optimizer(param_groups)
    self.scheduler = self._instantiate_scheduler()

    # ----------------------------
    # Precision
    # ----------------------------
    # Requested AMP follows the required priority:
    #   Trainer config explicit True/False wins.
    #   Trainer config None falls back to current().amp.
    #
    # Effective AMP is then hardened:
    #   - disabled when model caps.supports_amp is False
    #   - disabled when device does not support Trainer/DeepVisionModel autocast
    self.amp_requested = _resolve_requested_amp(cfg.precision.amp)
    self.amp_dtype = _amp_dtype(cfg.precision.amp_dtype, self.device)

    m = cast(Any, _unwrap(self.model))
    supports_amp = _model_supports_amp(m)
    device_supports_amp = _device_supports_trainer_amp(self.device)

    self.amp = bool(self.amp_requested and supports_amp and device_supports_amp)

    if hasattr(m, "use_amp"):
        m.use_amp = bool(self.amp)

    if hasattr(m, "autocast_dtype"):
        m.autocast_dtype = self.amp_dtype

    self.grad_clip = cfg.precision.grad_clip
    self.accumulate_steps = max(int(cfg.precision.accumulate_steps), 1)

    self.scaler = torch.amp.GradScaler(
        enabled=(
            self.amp
            and self.device.type == "cuda"
            and self.amp_dtype == torch.float16
        ),
        init_scale=2.0**16,
    )

    # ----------------------------
    # EMA
    # ----------------------------
    self.ema = None
    if cfg.ema.enabled:
        self.ema = ModelEMA(
            _unwrap(self.model),
            decay=cfg.ema.decay,
            warmup_steps=cfg.ema.warmup_steps,
        )

    # ----------------------------
    # State
    # ----------------------------
    self.global_step = 0
    self.epoch = 0
    self.epochs = int(cfg.epochs)
    self.should_stop = False
    self.metric_store: MutableMapping[str, Any] = {}

    # ----------------------------
    # Decide logging backends early
    # ----------------------------
    backends = self.cfg.logging.backends
    back_list = [backends] if isinstance(backends, str) else list(backends)

    use_mlflow = "mlflow" in back_list
    use_tensorboard = ("tensorboard" in back_list) or ("tb" in back_list)

    # ----------------------------
    # Resolve run / experiment naming
    # ----------------------------
    exp_name = self._resolve_effective_run_name()

    mlflow_args = dict(self.cfg.logging.mlflow or {})
    if getattr(self.cfg.logging, "mlflow_experiment_name", None):
        mlflow_args.setdefault(
            "experiment_name", str(self.cfg.logging.mlflow_experiment_name)
        )
    if getattr(self.cfg.logging, "mlflow_run_name", None):
        mlflow_args.setdefault("run_name", str(self.cfg.logging.mlflow_run_name))

    # ----------------------------
    # Logger(s) first (rank0 only)
    # ----------------------------
    if getattr(self.strategy, "is_rank0", True):
        self.logger = build_logger(
            backends=self.cfg.logging.backends,
            run_dir=Path("."),  # placeholder; actual run_dir decided below
            tb_args=self.cfg.logging.tensorboard,
            mlflow_args=mlflow_args,
            run_name=exp_name,
        )
    else:
        from .loggers import NullLogger

        self.logger = NullLogger()

    # ----------------------------
    # Run dir
    # ----------------------------
    run_dir = self._resolve_run_dir(
        exp_name=exp_name, use_mlflow=use_mlflow, use_tensorboard=use_tensorboard
    )
    run_dir.mkdir(parents=True, exist_ok=True)

    self.run_dir = run_dir
    self.checkpoint_manager = CheckpointManager(
        self.run_dir,
        enabled=bool(getattr(self.cfg.checkpoint, "enable_checkpointing", True)),
    )

    # ----------------------------
    # Export initial config snapshot as an artifact (rank0 only)
    # ----------------------------
    if getattr(self.strategy, "is_rank0", True):
        try:
            try:
                import yaml

                cfg_dict = self.cfg.export_dict(
                    portable=True, include_meta=True, use_initial=True
                )
                yaml_text = yaml.safe_dump(
                    cfg_dict,
                    sort_keys=False,
                    default_flow_style=False,
                    allow_unicode=True,
                )
            except Exception:
                cfg_dict = self.cfg.export_dict(
                    portable=True, include_meta=True, use_initial=True
                )
                yaml_text = str(cfg_dict)

            if hasattr(self.logger, "log_text_artifact"):
                self.logger.log_text_artifact(
                    yaml_text, artifact_file="config/initial_config.yaml"
                )
            else:
                cfg_path = self.run_dir / "initial_config.yaml"
                cfg_path.write_text(yaml_text, encoding="utf-8")
                with contextlib.suppress(Exception):
                    self.logger.log_artifact(cfg_path, artifact_path="config")
        except Exception:
            pass

    # ----------------------------
    # Log params snapshot
    # ----------------------------
    if getattr(self.strategy, "is_rank0", True):
        with contextlib.suppress(Exception):
            cfg_any = cast(Any, self.cfg)
            self.logger.log_params(
                {"config": cfg_any.to_dict(), "describe": cfg_any.describe()}
            )

    # ----------------------------
    # Hooks
    # ----------------------------
    self._hooks = self._build_default_hooks()
    self._metrics_monitor: Any = None

fit(resume=None)

Complete training program.

Parameters:

Name	Type	Description	Default
resume	Optional[Union[str, Path]], **optional**	Resume training from checkpoint. Defaults to None.	None

Source code in src/deepvisiontools/train/trainer.py

def fit(self, resume: str | Path | None = None) -> None:
    """Complete training program.

    Args:
        resume (Optional[Union[str, Path]], **optional**): Resume training from checkpoint. Defaults to None.
    """
    # Optionally resume
    if resume:
        self._resume_from(resume)

    # Re-apply any runtime normalization constraints before the first epoch
    self._reapply_norm_runtime_policy()

    # Rank-0 startup summary + artifact
    if getattr(self.strategy, "is_rank0", True):
        self._emit_norm_policy_summary()

    hooks: HookList = self._hooks
    hooks.before_train(self)
    try:
        for epoch in range(self.epoch, self.epochs):
            self.epoch = epoch
            self._metrics_monitor = MetricsMonitor(
                access=self.cfg.checkpoint.monitor
            )
            hooks.before_epoch(self, epoch)
            self._train_one_epoch()

            # Validation (EMA-aware)
            world_size = int(getattr(self.strategy, "world_size", 1))
            needs_eval = (
                self.valid_loader is not None or self._should_eval_train_val()
            )

            if world_size > 1 and needs_eval:
                self.validate_sync()
            else:
                self.validate()

            hooks.after_epoch(self, epoch, self._metrics_monitor)

            if self.should_stop:
                break
    except BaseException as exc:
        hooks.on_exception(self, exc)  # allow hooks to react
        raise
    finally:
        hooks.after_train(self)
        self.strategy.cleanup()
        with contextlib.suppress(Exception):
            self.logger.close()

validate()

Epoch-end validation.

DDP policy: - only rank0 runs the full validation pass - other ranks wait at barrier

Safety: - preserves the model's previous train/eval mode - restores BN runtime freeze policy when returning to train mode - stays symmetric with validate_sync()

Source code in src/deepvisiontools/train/trainer.py

@torch.no_grad()
def validate(self) -> Mapping[str, float]:
    """
    Epoch-end validation.

    DDP policy:
    - only rank0 runs the full validation pass
    - other ranks wait at barrier

    Safety:
    - preserves the model's previous train/eval mode
    - restores BN runtime freeze policy when returning to train mode
    - stays symmetric with validate_sync()
    """
    wants_train_val = self._should_eval_train_val()
    has_valid = self.valid_loader is not None

    if not has_valid and not wants_train_val:
        self.strategy.barrier()
        return {}

    if not getattr(self.strategy, "is_rank0", True):
        self.strategy.barrier()
        return {}

    valid_loader = self.valid_loader
    train_val_loader = self._build_train_val_loader() if wants_train_val else None

    if valid_loader is None and train_val_loader is None:
        self.strategy.barrier()
        return {}

    ema_hook = cast(Any, self._get_hook(EMAHook))
    cm = ema_hook.swapped(self) if ema_hook is not None else _null_cm()

    was_training = bool(self.model.training)

    try:
        out_dict: dict[str, float] = {}

        with cm:
            self.model.eval()

            if valid_loader is not None:
                out_dict.update(
                    self._run_epoch_eval_loader(
                        loader=valid_loader,
                        split="valid",
                        desc=f"Valid epoch {self.epoch}",
                    )
                )

            if train_val_loader is not None:
                out_dict.update(
                    self._run_epoch_eval_loader(
                        loader=train_val_loader,
                        split="train-val",
                        desc=f"Train-val epoch {self.epoch}",
                    )
                )

        monitor = getattr(self, "_metrics_monitor", None)
        if monitor is not None and isinstance(
            getattr(monitor, "metrics", None), dict
        ):
            for k, v in monitor.metrics.items():
                if isinstance(v, (int, float)):
                    out_dict[k] = float(v)

        return out_dict

    finally:
        if was_training:
            self.model.train()
            self._reapply_norm_runtime_policy()
        self.strategy.barrier()

export_model(path, *, builder, model_kwargs, include_optimizer=False, include_scaler=False, include_ema=True)

Write a portable model artifact (weights + builder/kwargs + pred params).

Source code in src/deepvisiontools/train/trainer.py

def export_model(
    self,
    path: str | Path,
    *,
    builder: str,
    model_kwargs: dict[str, Any],
    include_optimizer: bool = False,
    include_scaler: bool = False,
    include_ema: bool = True,
) -> None:
    """Write a portable model artifact (weights + builder/kwargs + pred params)."""
    model = _unwrap(self.model)
    optimizer = self.optimizer if include_optimizer else None
    scaler = self.scaler if (include_scaler and self.scaler is not None) else None
    ema_state = (
        self.ema.state_dict() if (include_ema and self.ema is not None) else None
    )

    save_checkpoint(
        str(path),
        model,
        builder=builder,
        model_kwargs=dict(model_kwargs),  # copy for safety
        epoch=self.epoch,
        optimizer=optimizer,  # optional
        scaler=cast(Any, scaler),  # optional
        ema_state_dict=ema_state,  # optional
        extra={"global_step": self.global_step},
    )

validate_sync()

Synchronized cross-rank in-epoch validation.

Design: - all ranks participate - validation data is sharded across ranks - uses a temporary fresh valid meter from MetricsHook - final metrics are synchronized through suite.finalize_distributed() - rank0 updates the monitor + logger - restores previous train/eval mode afterwards

Source code in src/deepvisiontools/train/trainer.py

@torch.no_grad()
def validate_sync(self) -> Mapping[str, float]:
    """
    Synchronized cross-rank in-epoch validation.

    Design:
    - all ranks participate
    - validation data is sharded across ranks
    - uses a temporary fresh valid meter from MetricsHook
    - final metrics are synchronized through suite.finalize_distributed()
    - rank0 updates the monitor + logger
    - restores previous train/eval mode afterwards
    """
    eval_jobs: list[tuple[Literal["valid", "train-val"], Any]] = []

    world_size = int(getattr(self.strategy, "world_size", 1))
    if world_size <= 1:
        return self.validate()

    wants_train_val = self._should_eval_train_val()

    if self.valid_loader is not None:
        valid_loader = self._build_rank_sharded_eval_loader(
            self.valid_loader,
            batch_augmenter=None,
        )
        if valid_loader is not None:
            eval_jobs.append(("valid", valid_loader))

    if wants_train_val:
        train_val_loader = self._build_train_val_loader()
        if train_val_loader is not None:
            train_val_loader = self._build_rank_sharded_eval_loader(
                train_val_loader,
                batch_augmenter=None,
            )
            if train_val_loader is not None:
                eval_jobs.append(("train-val", train_val_loader))

    if not eval_jobs:
        self.strategy.barrier()
        return {}

    rank0 = bool(getattr(self.strategy, "is_rank0", True))
    was_training = bool(self.model.training)

    ema_hook = cast(Any, self._get_hook(EMAHook))
    metrics_hook = cast(Any, self._get_hook(MetricsHook))
    cm = ema_hook.swapped(self) if ema_hook is not None else _null_cm()

    loss_meters: dict[Literal["valid", "train-val"], LossAverager] = {
        job_split: LossAverager() for job_split, _loader in eval_jobs
    }

    temp_meters: dict[Literal["valid", "train-val"], Any] = {}
    if metrics_hook is not None:
        for job_split, _loader in eval_jobs:
            try:
                if job_split == "valid" and hasattr(
                    metrics_hook,
                    "new_valid_meter",
                ):
                    temp_meters[job_split] = metrics_hook.new_valid_meter()
                elif job_split == "train-val" and hasattr(
                    metrics_hook,
                    "new_train_meter",
                ):
                    temp_meters[job_split] = metrics_hook.new_train_meter()
                else:
                    temp_meters[job_split] = None
            except Exception:
                temp_meters[job_split] = None

    previous_ddp_eval_forward = bool(
        getattr(self, "_ddp_eval_forward_enabled", False)
    )

    def _reduce_loss_meter(
        meter: LossAverager,
    ) -> dict[str, float]:
        local_count = torch.tensor(
            [float(meter.count)],
            dtype=torch.float64,
            device=self.device,
        )

        local_keys = sorted(str(k) for k in meter.total)
        global_keys = list(local_keys)

        if dist.is_available() and dist.is_initialized():
            gathered_keys: list[list[str]] = [[] for _ in range(world_size)]
            dist.all_gather_object(gathered_keys, local_keys)
            global_keys = sorted({k for keys in gathered_keys for k in keys})

        local_sums = torch.tensor(
            [float(meter.total.get(k, 0.0)) for k in global_keys],
            dtype=torch.float64,
            device=self.device,
        )

        if dist.is_available() and dist.is_initialized():
            dist.all_reduce(local_count, op=dist.ReduceOp.SUM)
            if local_sums.numel() > 0:
                dist.all_reduce(local_sums, op=dist.ReduceOp.SUM)

        total_count = float(local_count.item())
        if total_count <= 0.0 or not global_keys:
            return {}

        return {
            k: float(local_sums[i].item()) / total_count
            for i, k in enumerate(global_keys)
        }

    try:
        self._ddp_eval_forward_enabled = True

        with cm:
            self.model.eval()

            for job_split, eval_loader in eval_jobs:
                loss_meter = loss_meters[job_split]

                iterator = (
                    tqdm(
                        enumerate(eval_loader),
                        total=len(eval_loader),
                        desc=f"{job_split}(sync) epoch {self.epoch}",
                    )
                    if rank0
                    else enumerate(eval_loader)
                )

                for _step, batch in iterator:
                    with torch.autocast(
                        self.device.type,
                        enabled=self.amp,
                        dtype=self.amp_dtype,
                    ):
                        out: ModelOutput = self._call_eval_step(batch)
                        loss = out.globalloss() if out.loss is not None else None

                    if loss is not None:
                        with contextlib.suppress(Exception):
                            loss_meter.update(out.lossesdict())

                    temp_meter = temp_meters.get(job_split)
                    if temp_meter is not None:
                        with contextlib.suppress(Exception):
                            _imgs, targets, _extra_args = self._unpack_batch(batch)
                            if out.preds is not None and targets is not None:
                                temp_meter.update(out.preds, targets)

                    if rank0 and hasattr(iterator, "set_postfix_str"):
                        iterator.set_postfix_str(
                            _log_string(loss_meter.avg)
                            if loss_meter.count > 0
                            else ""
                        )

        out_dict: dict[str, float] = {}

        # ---- reduce losses across ranks, split by split ----
        for meter_split, loss_meter in loss_meters.items():
            avg = _reduce_loss_meter(loss_meter)
            if not avg:
                continue

            if rank0 and getattr(self, "_metrics_monitor", None) is not None:
                self._metrics_monitor.losses.update({f"{meter_split}/loss": avg})

            if "loss" in avg:
                out_dict[f"{meter_split}/loss"] = float(avg["loss"])

            sync_prefix = (
                "valid_sync" if meter_split == "valid" else "train-val_sync"
            )
            for loss_name, loss_value in avg.items():
                out_dict[f"{sync_prefix}/{loss_name}"] = float(loss_value)

        # ---- distributed metrics suite finalize, split by split ----
        for meter_split, temp_meter in temp_meters.items():
            if temp_meter is None:
                continue

            metric_results: dict[str, float] = {}
            try:
                if hasattr(temp_meter, "finalize_distributed"):
                    metric_results = temp_meter.finalize_distributed()
                else:
                    metric_results = temp_meter.finalize()
            except Exception:
                metric_results = {}

            if (
                rank0
                and metric_results
                and getattr(self, "_metrics_monitor", None) is not None
            ):
                prefixed = {
                    f"{meter_split}/{k}": float(v)
                    for k, v in metric_results.items()
                }
                self._metrics_monitor.metrics.update(prefixed)
                out_dict.update(prefixed)

        if rank0 and out_dict:
            with contextlib.suppress(Exception):
                self.logger.log_metrics(
                    out_dict,
                    step=int(self.global_step),
                    split="step",
                )

        return out_dict if rank0 else {}

    finally:
        self._ddp_eval_forward_enabled = previous_ddp_eval_forward

        if was_training:
            self.model.train()
            self._reapply_norm_runtime_policy()

        self.strategy.barrier()

launch(main_fn, *, devices='auto', nproc=None, main_args=None, main_kwargs=None)

Launch a training entrypoint, optionally through torch.distributed.run.

launch is a convenience wrapper around torchrun --standalone. If the current process is already inside a torchrun context (LOCAL_RANK is set), it directly calls main_fn(*main_args, **main_kwargs). Otherwise, it infers the number of processes from devices or nproc.

Parameters:

Name	Type	Description	Default
main_fn	`Callable[..., Any]`	User training function to execute.	required
devices	`str | int | Sequence[int]`	Device/process selection. Accepted values include: - "auto": use torch.cuda.device_count() processes, or 1 if none. - "cpu" or "single": run one in-process training function. - int: process count, not a physical GPU id. - [0, 1, 3]: expose exactly those physical GPU ids to the child. - "0,1,3" or "cuda:0,1,3": same GPU-subset behavior. - "cuda:2": expose physical GPU 2 as a single visible GPU in a child process.	'auto'
nproc	`int | None`	Optional explicit process count override. If set, it takes precedence over the process count inferred from devices.	None
main_args	`Sequence[Any] | None`	Positional arguments forwarded to main_fn.	None
main_kwargs	`Mapping[str, Any] | None`	Keyword arguments forwarded to main_fn.	None

DDP behavior

• nproc == 1 and no GPU remapping: main_fn runs in-process.
• nproc == 1 with GPU remapping, for example devices="cuda:2": a child process is spawned so PyTorch sees the intended CUDA_VISIBLE_DEVICES.
• nproc > 1: the current Python script is re-executed with python -m torch.distributed.run --standalone --nproc_per_node=<nproc>.

Usage

from deepvisiontools.train import Trainer, launch

def main():
    Trainer(cfg).fit()

# Auto: one process per visible GPU, or CPU/single-process if no GPU.
launch(main, devices="auto")

# Force CPU / single process.
launch(main, devices="cpu")

# Use 2 processes. This is a process count, not GPU id 2.
launch(main, devices=2)

# Use exactly physical GPUs 0 and 1.
launch(main, devices=[0, 1])

# Same as above with string syntax.
launch(main, devices="cuda:0,1")

# Use only physical GPU 2.
launch(main, devices="cuda:2")

# Forward arguments to the main function.
def train_with_name(run_name, *, resume=None):
    cfg.logging.mlflow_run_name = run_name
    Trainer(cfg).fit(resume=resume)

launch(
    train_with_name,
    devices="auto",
    main_args=("experiment_001",),
    main_kwargs={"resume": None},
)

Source code in src/deepvisiontools/train/trainer.py

def launch(
    main_fn,
    *,
    devices: str | int | Sequence[int] = "auto",
    nproc: int | None = None,
    main_args: Sequence[Any] | None = None,
    main_kwargs: Mapping[str, Any] | None = None,
) -> None:
    """
    Launch a training entrypoint, optionally through `torch.distributed.run`.

    `launch` is a convenience wrapper around `torchrun --standalone`. If the
    current process is already inside a torchrun context (`LOCAL_RANK` is set),
    it directly calls `main_fn(*main_args, **main_kwargs)`. Otherwise, it infers
    the number of processes from `devices` or `nproc`.

    Args:
        main_fn (`Callable[..., Any]`): User training function to execute.
        devices (`str | int | Sequence[int]`): Device/process selection.
            Accepted values include:
            - `"auto"`: use `torch.cuda.device_count()` processes, or 1 if none.
            - `"cpu"` or `"single"`: run one in-process training function.
            - `int`: process count, not a physical GPU id.
            - `[0, 1, 3]`: expose exactly those physical GPU ids to the child.
            - `"0,1,3"` or `"cuda:0,1,3"`: same GPU-subset behavior.
            - `"cuda:2"`: expose physical GPU 2 as a single visible GPU in a
              child process.
        nproc (`int | None`): Optional explicit process count override. If set,
            it takes precedence over the process count inferred from `devices`.
        main_args (`Sequence[Any] | None`): Positional arguments forwarded to
            `main_fn`.
        main_kwargs (`Mapping[str, Any] | None`): Keyword arguments forwarded to
            `main_fn`.

    ???+ note "DDP behavior"
        - `nproc == 1` and no GPU remapping: `main_fn` runs in-process.
        - `nproc == 1` with GPU remapping, for example `devices="cuda:2"`:
          a child process is spawned so PyTorch sees the intended
          `CUDA_VISIBLE_DEVICES`.
        - `nproc > 1`: the current Python script is re-executed with
          `python -m torch.distributed.run --standalone --nproc_per_node=<nproc>`.

    ???+ example "Usage"
        ```python
        from deepvisiontools.train import Trainer, launch

        def main():
            Trainer(cfg).fit()

        # Auto: one process per visible GPU, or CPU/single-process if no GPU.
        launch(main, devices="auto")

        # Force CPU / single process.
        launch(main, devices="cpu")

        # Use 2 processes. This is a process count, not GPU id 2.
        launch(main, devices=2)

        # Use exactly physical GPUs 0 and 1.
        launch(main, devices=[0, 1])

        # Same as above with string syntax.
        launch(main, devices="cuda:0,1")

        # Use only physical GPU 2.
        launch(main, devices="cuda:2")

        # Forward arguments to the main function.
        def train_with_name(run_name, *, resume=None):
            cfg.logging.mlflow_run_name = run_name
            Trainer(cfg).fit(resume=resume)

        launch(
            train_with_name,
            devices="auto",
            main_args=("experiment_001",),
            main_kwargs={"resume": None},
        )
        ```
    """
    if os.environ.get("LOCAL_RANK") is not None:
        return main_fn(*(main_args or []), **(main_kwargs or {}))

    visible_subset: str | None = None

    def _parse_device_string(s: str) -> tuple[int, str | None]:
        raw = str(s).strip().lower()

        if raw in {"auto", "cuda", "gpu", "gpus"}:
            return max(1, get_visible_gpu_count()), None

        if raw in {"cpu", "single"}:
            return 1, None

        had_cuda_prefix = raw.startswith("cuda:")
        if had_cuda_prefix:
            raw = raw[len("cuda:") :]

        if "," in raw:
            gpu_ids = [tok.strip() for tok in raw.split(",") if tok.strip()]
            if gpu_ids and all(tok.isdigit() for tok in gpu_ids):
                return len(gpu_ids), ",".join(gpu_ids)

        if raw.isdigit():
            if had_cuda_prefix:
                # "cuda:2" means: expose physical GPU 2 as the only visible GPU.
                return 1, raw

            # Historical behavior:
            # "2" means process count, not physical GPU id 2.
            # Use "cuda:2" to select a physical single-GPU subset.
            return int(raw), None

        raise ValueError(
            "launch(devices=...) must be 'auto', 'cpu', an int, a sequence of GPU ids, "
            "'cuda:N', 'N,M', or 'cuda:N,M'. "
            f"Got {s!r}."
        )

    if nproc is None:
        if devices == "auto":
            nproc = max(1, get_visible_gpu_count())

        elif isinstance(devices, int):
            nproc = int(devices)

        elif isinstance(devices, (list, tuple)):
            gpu_ids = [str(int(d)) for d in devices]
            nproc = len(gpu_ids)
            if gpu_ids:
                visible_subset = ",".join(gpu_ids)

        elif isinstance(devices, str):
            nproc, visible_subset = _parse_device_string(devices)

        else:
            nproc = 1

    nproc = max(1, int(nproc))

    env = os.environ.copy()
    if visible_subset is not None:
        env["CUDA_VISIBLE_DEVICES"] = visible_subset

    if nproc == 1:
        # For single-GPU subset selection, run the user function in-process only
        # when no CUDA_VISIBLE_DEVICES remapping is needed. If remapping is needed,
        # re-exec a child process so torch sees the intended visible subset.
        if visible_subset is None:
            return main_fn(*(main_args or []), **(main_kwargs or {}))

        cmd = [sys.executable, sys.argv[0], *sys.argv[1:]]
        subprocess.run(cmd, check=True, env=env)
        return

    cmd = [
        sys.executable,
        "-m",
        "torch.distributed.run",
        "--standalone",
        f"--nproc_per_node={nproc}",
        sys.argv[0],
        *sys.argv[1:],
    ]

    subprocess.run(cmd, check=True, env=env)

TrainerConfig(model=None, train_loader=None, valid_loader=None, model_spec=None, epochs=100, output_dir=None, devices=DevicesConfig(), optim=OptimConfig(), sched=SchedConfig(), precision=PrecisionConfig(), ema=EMAConfig(), freezing=FreezingConfig(), normalization=NormalizationConfig(), checkpoint=CheckpointConfig(), logging=LoggingConfig(), eval=EvalConfig(), find_unused_parameters=False, optuna=None) dataclass

Main configuration object consumed by Trainer.

TrainerConfig gathers runtime objects such as loaders, portable model reconstruction information, training duration, device/DDP policy, optimization, scheduler/warmup, precision, EMA, freezing, normalization, checkpointing, logging, evaluation, and optional Optuna configuration.

Typical usage is to start from a registered preset, then override only what is specific to the experiment.

Parameters:

Name	Type	Description	Default
model	`Any | None`	Optional runtime model object. If provided, it is wrapped into ModelSpec when possible, then stored under model_spec. Prefer model_spec for portable YAML/export workflows.	None
train_loader	`Any`	Training dataloader. A DeepVisionLoader is recommended.	None
valid_loader	`Any`	Validation dataloader. A DeepVisionLoader is recommended. Can be None for train-only workflows.	None
model_spec	`ModelSpec | None`	Portable model reconstruction information. It stores the model builder key and kwargs and can also hold the runtime model object.	None
epochs	`int`	Number of training epochs. Defaults to 100.	100
output_dir	`str | None`	Optional run output directory. If None, deepvisiontools uses its default output/logging setup.	None
devices	`DevicesConfig`	Device, strategy, num_workers, and pin_memory policy.	deepvisiontools.train.trainconfig.DevicesConfig()
optim	`OptimConfig`	Optimizer class/spec, optimizer kwargs, weight decay, and parameter-group strategy.	deepvisiontools.train.trainconfig.OptimConfig()
sched	`SchedConfig`	Scheduler class/spec, scheduler kwargs, step policy, and optional warmup.	deepvisiontools.train.trainconfig.SchedConfig()
precision	`PrecisionConfig`	AMP, AMP dtype, gradient clipping, and gradient accumulation.	deepvisiontools.train.trainconfig.PrecisionConfig()
ema	`EMAConfig`	Exponential moving average configuration.	deepvisiontools.train.trainconfig.EMAConfig()
freezing	`FreezingConfig`	Optional temporary backbone freezing.	deepvisiontools.train.trainconfig.FreezingConfig()
normalization	`NormalizationConfig`	BatchNorm / SyncBatchNorm policy.	deepvisiontools.train.trainconfig.NormalizationConfig()
checkpoint	`CheckpointConfig`	Training checkpoint and portable model artifact export configuration.	deepvisiontools.train.trainconfig.CheckpointConfig()
logging	`LoggingConfig`	Logging backend configuration, including MLflow and TensorBoard settings.	deepvisiontools.train.trainconfig.LoggingConfig()
eval	`EvalConfig`	Validation frequency, evaluated split(s), and metric suite configuration.	deepvisiontools.train.trainconfig.EvalConfig()
find_unused_parameters	`bool`	DDP option forwarded to DistributedDataParallel. Useful for models with conditionally unused branches, but it can add overhead. Defaults to False.	False
optuna	`Any | None`	Optional Optuna configuration. In normal training it can stay None. For HPO, set it to an OptunaConfig and call run_optuna(cfg).	None

Recommended usage from a registered preset

from deepvisiontools.train import TrainerConfig

cfg = TrainerConfig.build_registered_config(
    model="detector-seg",
    config="detector_convnextv2-base_cascade_maskrcnn",
    train_loader=trainloader,
    valid_loader=validloader,
)

cfg.epochs = 50
cfg.devices.device = "cuda"
cfg.devices.strategy = "auto"
cfg.devices.num_workers = 4
cfg.devices.pin_memory = True

cfg.precision.amp = True
cfg.precision.amp_dtype = "bf16"
cfg.precision.grad_clip = 5.0
cfg.precision.accumulate_steps = 2

cfg.ema.enabled = True
cfg.ema.decay = 0.995
cfg.ema.update_every = 1
cfg.ema.warmup_steps = 100
cfg.ema.eval_with_ema = True

cfg.checkpoint.enable_checkpointing = True
cfg.checkpoint.monitor = ("valid/loss", "loss")
cfg.checkpoint.mode = "min"
cfg.checkpoint.filename = "checkpoint"
cfg.checkpoint.save_every_n_epochs = 1
cfg.checkpoint.keep_last_n = 3
cfg.checkpoint.export_model_artifact = False
cfg.checkpoint.export_when = "best"

cfg.logging.backends = ["mlflow"]
cfg.logging.mlflow_experiment_name = "deepvisiontools_experiments"
cfg.logging.mlflow_run_name = "detector_seg_baseline"
cfg.logging.log_every_n_steps = 0

cfg.eval.eval_every_n_steps = 0
cfg.eval.eval_on = "valid"

Optimizer, scheduler, warmup and LLRD

import torch

from deepvisiontools.train import TrainerConfig
from deepvisiontools.train.trainconfig import WarmupConfig

cfg = TrainerConfig.build_registered_config(
    model="detector",
    config="detector_convnextv2-base_vfnet",
    train_loader=trainloader,
    valid_loader=validloader,
)

cfg.optim.optimizer = torch.optim.AdamW
cfg.optim.optim_args = {
    "lr": 2.5e-4,
    "betas": (0.9, 0.999),
}
cfg.optim.weight_decay = 0.05

# Available built-in modes are "flat", "llrd", and "layer_aware_llrd".
cfg.optim.param_groups = "layer_aware_llrd"
cfg.optim.layer_decay = 0.85

cfg.sched.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR
cfg.sched.scheduler_args = {"T_max": "auto", "eta_min": 1e-6}
cfg.sched.step_on = "step"
cfg.sched.warmup = WarmupConfig(
    steps=500,
    steps_unit="step",
    ratio=1.0,
    start_factor=0.001,
)

Evaluation metrics

from deepvisiontools.train.trainconfig import EvalConfig, MetricsSpec

# Detection metrics.
cfg.eval = EvalConfig(
    eval_every_n_steps=0,
    eval_on="valid",
    metrics=MetricsSpec(
        kind="bbox",
        num_classes=3,
    ),
)

# Instance segmentation metrics.
cfg.eval.metrics = MetricsSpec(
    kind="instance_mask",
    num_classes=3,
)

# Semantic segmentation metrics.
cfg.eval.metrics = MetricsSpec(
    kind="semantic_mask",
    num_classes=3,
    kwargs={"keep_per_image": True},
)

Portable custom model with ModelSpec

# my_project/custom_models.py
#
# The custom builder must be imported before loading/building the config,
# so that it is registered in MODEL_REGISTRY.
#
# IMPORTANT:
# - the builder must return a DeepVisionModel;
# - all kwargs stored in ModelSpec must be YAML-safe;
# - avoid storing runtime objects such as instantiated torch modules,
#   datasets, losses or lambdas in model_spec.kwargs.

import torch

from deepvisiontools.models.base.basemodel import (
    MODEL_REGISTRY,
    DeepVisionModel,
)

@MODEL_REGISTRY.register("my_custom_smp")
def build_my_custom_smp(
    *,
    arch: str = "Unet",
    encoder_name: str = "resnet18",
    encoder_weights: str | None = "imagenet",
    in_channels: int = 3,
    num_classes: int = 1,
    device: str | torch.device | None = None,
    use_amp: bool | None = None,
) -> DeepVisionModel:
    from deepvisiontools.models import ModelFactory

    model = ModelFactory(
        name="smp",
        arch=arch,
        encoder_name=encoder_name,
        encoder_weights=encoder_weights,
        in_channels=in_channels,
        num_classes=num_classes,
        device=device,
        use_amp=use_amp,
    )

    # Ensure exported artifacts and TrainerConfig YAML can rebuild this
    # exact model through the registered custom builder.
    model.set_export_spec(
        "my_custom_smp",
        {
            "arch": arch,
            "encoder_name": encoder_name,
            # Usually set to None for artifact reload because weights
            # will be loaded from the saved state_dict.
            "encoder_weights": None,
            "in_channels": in_channels,
            "num_classes": num_classes,
            "device": device,
            "use_amp": use_amp,
        },
    )

    return model

# train.py

# Import the module once so @MODEL_REGISTRY.register(...) is executed.
import my_project.custom_models  # noqa: F401

import torch

from deepvisiontools.train import Trainer, TrainerConfig
from deepvisiontools.train.trainconfig import (
    ModelSpec,
    OptimConfig,
    SchedConfig,
)

cfg = TrainerConfig(
    train_loader=trainloader,
    valid_loader=validloader,
    model_spec=ModelSpec(
        builder="my_custom_smp",
        kwargs={
            "arch": "Unet",
            "encoder_name": "resnet18",
            "encoder_weights": "imagenet",
            "in_channels": 3,
            "num_classes": 1,
            "device": "cuda",
            "use_amp": True,
        },
    ),
    epochs=30,
)

cfg.optim = OptimConfig(
    optimizer=torch.optim.AdamW,
    optim_args={"lr": 1e-4},
    weight_decay=1e-4,
)

cfg.sched = SchedConfig(
    scheduler=None,
    scheduler_args={},
    step_on="epoch",
)

Trainer(cfg).fit()

The exported YAML will contain a portable model block similar to:

trainer:
  model_spec:
    builder: my_custom_smp
    kwargs:
      arch: Unet
      encoder_name: resnet18
      encoder_weights: imagenet
      in_channels: 3
      num_classes: 1
      device: cuda
      use_amp: true

To reload this YAML later, import the custom builder module before calling TrainerConfig.from_yaml(...):

import my_project.custom_models  # noqa: F401

from deepvisiontools.train import TrainerConfig

cfg = TrainerConfig.from_yaml(
    "my_training_config.yaml",
    train_loader=trainloader,
    valid_loader=validloader,
)

Normalization, freezing and DDP

# Device strategy:
# - "auto": single process unless launched with torchrun/launch(...).
# - "single": force single-device training.
# - "ddp": require distributed context.
cfg.devices.device = "cuda"
cfg.devices.strategy = "auto"

# Required by some DDP models with conditional branches.
cfg.find_unused_parameters = False

# BatchNorm policy.
cfg.normalization.policy = "auto"
cfg.normalization.freeze_affine = True
cfg.normalization.freeze_running_stats = True
cfg.normalization.verbose = True

# Temporarily freeze backbone optimizer param groups.
cfg.freezing.freeze_epochs = 2
cfg.freezing.backbone_group_name = "backbone"

Optuna HPO entrypoint

from deepvisiontools.train.optuna_wrapper import (
    OptunaConfig,
    SearchSpaceSpec,
    run_optuna,
)

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="my_study",
    storage="my_study.db",
    n_trials=20,
    search_space=[
        SearchSpaceSpec(
            name="lr",
            kind="float",
            path="optim.optim_args.lr",
            low=1e-5,
            high=1e-3,
            log=True,
        ),
        SearchSpaceSpec(
            name="accumulate_steps",
            kind="categorical",
            path="precision.accumulate_steps",
            choices=[1, 2, 4],
        ),
    ],
)

study = run_optuna(cfg)

ensure_model_spec_from_model()

If a DeepVisionModel exists in model_spec and has export_spec(), populate model_spec/checkpoint hints.

Source code in src/deepvisiontools/train/trainconfig.py

def ensure_model_spec_from_model(self) -> None:
    """If a DeepVisionModel exists in model_spec and has export_spec(), populate model_spec/checkpoint hints."""
    # In current design, runtime model lives under model_spec._model
    if self.model_spec is None:
        return

    m = getattr(self.model_spec, "model_obj", None)
    if m is None:
        m = getattr(self.model_spec, "_model", None)

    if m is None or not isinstance(m, DeepVisionModel):
        return

    try:
        spec = m.export_spec()
    except Exception:
        spec = None

    if not (isinstance(spec, tuple) and len(spec) == 2):
        return

    builder, kwargs = spec
    builder = str(builder)
    kwargs = dict(kwargs)

    # Fill model_spec portable fieldswandb_project': None, 'wandb_run_name': N
    if not self.model_spec.builder:
        self.model_spec.builder = builder
    if not self.model_spec.kwargs:
        self.model_spec.kwargs = dict(kwargs)

    # Mirror into checkpoint hints
    if self.checkpoint.model_builder is None:
        self.checkpoint.model_builder = builder
    if self.checkpoint.model_kwargs is None:
        self.checkpoint.model_kwargs = dict(kwargs)

freeze_initial_snapshot()

Freeze a stable snapshot for reproducible export.

Source code in src/deepvisiontools/train/trainconfig.py

def freeze_initial_snapshot(self) -> None:
    """Freeze a stable snapshot for reproducible export."""
    # Make sure we capture model_spec if possible before freezing
    self.ensure_model_spec_from_model()
    self._initial_snapshot = {
        "portable": self._export_dict(portable=True, include_meta=False),
        "record": self._export_dict(portable=False, include_meta=True),
    }
    self._initial_frozen = True

from_yaml(path, *, train_loader, valid_loader=None, model=None) classmethod

Load portable YAML and rebuild a TrainerConfig.

You MUST pass loaders. You MAY pass model to override; otherwise if model_spec exists in YAML we rebuild from it.

Source code in src/deepvisiontools/train/trainconfig.py

@classmethod
def from_yaml(
    cls,
    path: str | Path,
    *,
    train_loader: Any,
    valid_loader: Any = None,
    model: Any = None,
) -> TrainerConfig:
    """Load portable YAML and rebuild a TrainerConfig.

    You MUST pass loaders. You MAY pass model to override; otherwise if model_spec
    exists in YAML we rebuild from it.
    """
    if yaml is None:  # pragma: no cover
        raise RuntimeError("PyYAML is not available; cannot import YAML.")
    raw_text = _read_yaml_source_text(path)
    data = yaml.safe_load(raw_text) or {}
    trainer = data.get("trainer", data)

    # model spec (optional)
    model_spec = None
    ms = trainer.get("model_spec", None)
    if isinstance(ms, dict) and "builder" in ms:
        model_spec = ModelSpec(
            builder=str(ms["builder"]), kwargs=dict(ms.get("kwargs") or {})
        )

    def _dc(dc_cls: Any, payload: Any) -> Any:
        return dc_cls(**payload) if isinstance(payload, dict) else dc_cls()

    # optional optuna config (lazy import to avoid circular import)
    optuna_cfg = None
    opt_payload = trainer.get("optuna", None)
    if isinstance(opt_payload, dict):
        try:
            from .optuna_wrapper import OptunaConfig

            optuna_cfg = OptunaConfig.from_dict(opt_payload)
        except Exception as e:
            raise ValueError(f"Failed to parse trainer.optuna: {e}") from e

    cfg = cls(
        model=model,
        model_spec=model_spec,
        train_loader=train_loader,
        valid_loader=valid_loader,
        epochs=int(trainer.get("epochs", 100)),
        output_dir=trainer.get("output_dir"),
        devices=_dc(DevicesConfig, trainer.get("devices")),
        precision=_dc(PrecisionConfig, trainer.get("precision")),
        ema=_dc(EMAConfig, trainer.get("ema")),
        freezing=_dc(FreezingConfig, trainer.get("freezing")),
        normalization=_dc(NormalizationConfig, trainer.get("normalization")),
        checkpoint=_dc(CheckpointConfig, trainer.get("checkpoint")),
        logging=_dc(LoggingConfig, trainer.get("logging")),
        find_unused_parameters=bool(trainer.get("find_unused_parameters", False)),
        optuna=optuna_cfg,
    )

    # optim
    opt_payload = trainer.get("optim") or {}
    opt_spec = opt_payload.get("optimizer")
    if not isinstance(opt_spec, dict) or "target" not in opt_spec:
        raise ValueError("optim.optimizer must be present with {target, kwargs}.")
    cfg.optim.optimizer = ImportSpec(
        target=str(opt_spec["target"]), kwargs=dict(opt_spec.get("kwargs") or {})
    )
    cfg.optim.optim_args = dict(
        opt_payload.get("optim_args") or cfg.optim.optimizer.kwargs
    )
    cfg.optim.param_groups = opt_payload.get("param_groups", "flat")
    cfg.optim.weight_decay = opt_payload.get("weight_decay")
    cfg.optim.layer_decay = opt_payload.get("layer_decay")

    # sched
    sch_payload = trainer.get("sched") or {}
    sch_spec = sch_payload.get("scheduler")
    if sch_spec is None:
        cfg.sched.scheduler = None
    else:
        if not isinstance(sch_spec, dict) or "target" not in sch_spec:
            raise ValueError("sched.scheduler must be null or {target, kwargs}.")
        cfg.sched.scheduler = ImportSpec(
            target=str(sch_spec["target"]),
            kwargs=dict(sch_spec.get("kwargs") or {}),
        )
    cfg.sched.scheduler_args = dict(
        sch_payload.get("scheduler_args")
        or (cfg.sched.scheduler.kwargs if cfg.sched.scheduler else {})
    )
    cfg.sched.step_on = sch_payload.get("step_on", "epoch")
    wu = sch_payload.get("warmup")
    cfg.sched.warmup = WarmupConfig(**wu) if isinstance(wu, dict) else None

    # eval
    ev_payload = trainer.get("eval") or {}
    metrics_spec = None
    m = ev_payload.get("metrics")
    if m is not None:
        if not isinstance(m, dict):
            raise ValueError("eval.metrics must be null or a dict.")
        metrics_spec = MetricsSpec(
            kind=m.get("kind"),
            num_classes=m.get("num_classes"),
            target=m.get("target"),
            kwargs=dict(m.get("kwargs") or {}),
        )

    cfg.eval = EvalConfig(
        eval_every_n_steps=int(ev_payload.get("eval_every_n_steps", 0)),
        eval_on=ev_payload.get("eval_on", "valid"),
        metrics=metrics_spec,
    )

    # Build model if not provided but spec exists
    if cfg.model is None and cfg.model_spec is not None:
        cfg.build_model()

    return cfg

available_configs() classmethod

List all registered default TrainerConfig presets.

Layout matches ModelFactory configs

dtype: model: - "name": description

Source code in src/deepvisiontools/train/trainconfig.py

@classmethod
def available_configs(cls) -> AvailableConfigs:
    """List all registered default TrainerConfig presets.

    Layout matches ModelFactory configs:
      dtype:
        model:
          - "name": description
    """
    _ensure_trainer_config_presets_loaded()

    tree: dict[str, dict[str, list[ConfigEntry]]] = {}
    for (dtype, model, name), preset in _TRAINER_CONFIG_PRESETS.items():
        tree.setdefault(dtype, {})
        tree[dtype].setdefault(model, [])
        tree[dtype][model].append(
            ConfigEntry(name=str(name), description=str(preset.description or ""))
        )
    return AvailableConfigs(tree)

build_registered_config(*, model, config, train_loader, valid_loader=None, dtype=None, model_override=None) classmethod

Build a TrainerConfig from a registered preset.

???+ example "Usage":

# example usage
Config = TrainerConfig.build_registered_config(
        model="detector-seg",
        config="detector_convnextv2-base_cascade_maskrcnn",
        train_loader=trainloader,
        valid_loader=validloader,
    )

Source code in src/deepvisiontools/train/trainconfig.py

@classmethod
def build_registered_config(
    cls,
    *,
    model: str,
    config: str,
    train_loader: Any,
    valid_loader: Any = None,
    dtype: str | None = None,
    model_override: Any = None,
) -> TrainerConfig:
    """Build a TrainerConfig from a registered preset.

    ???+ example "Usage":
        ```python
        # example usage
        Config = TrainerConfig.build_registered_config(
                model="detector-seg",
                config="detector_convnextv2-base_cascade_maskrcnn",
                train_loader=trainloader,
                valid_loader=validloader,
            )
        ```
    """
    _ensure_trainer_config_presets_loaded()

    model = str(model)
    config = str(config)

    if dtype is not None:
        dtype = str(dtype)
        preset = _TRAINER_CONFIG_PRESETS.get((dtype, model, config))
        if preset is None:
            raise KeyError(
                f"Unknown TrainerConfig preset: dtype={dtype} model={model} config={config}"
            )
    else:
        hits = [
            p
            for (dt, m, n), p in _TRAINER_CONFIG_PRESETS.items()
            if (m == model and n == config)
        ]
        if len(hits) == 0:
            raise KeyError(
                f"Unknown TrainerConfig preset: model={model} config={config}"
            )
        if len(hits) > 1:
            choices = ", ".join(sorted({p.dtype for p in hits}))
            raise KeyError(
                f"Ambiguous TrainerConfig preset for model={model} config={config}. "
                f"Please specify dtype (choices: {choices})."
            )
        preset = hits[0]

    yaml_text = _read_yaml_source_text(preset.source)
    return cls.from_yaml(
        yaml_text,
        train_loader=train_loader,
        valid_loader=valid_loader,
        model=model_override,
    )

ModelSpec(builder=None, kwargs=None, model=None) dataclass

Portable model reconstruction information + optional cached runtime model.

Note

• builder/kwargs are the portable representation (what goes to YAML)
• model is a runtime-only object (never exported in portable=True)
• model is a property so that setting it auto-refreshes builder/kwargs from DeepVisionModel.export_spec() when possible.
• If model is None at init: build from builder/kwargs (must be set).
• If model is provided at init and it's DeepVisionModel: override builder/kwargs from model.export_spec() when available.

Source code in src/deepvisiontools/train/trainconfig.py

def __init__(
    self,
    builder: str | None = None,
    kwargs: dict[str, Any] | None = None,
    model: Any | None = None,
) -> None:
    self.builder = str(builder) if builder is not None else None
    self.kwargs = dict(kwargs or {})
    self._model = None

    if model is not None:
        self.model = model
    elif self.builder:
        self._ensure_built_from_spec()

model_obj property

Internal getter with a clear name.

ImportSpec(target, kwargs=dict()) dataclass

Main class to make config exportable/readable from yaml.

DevicesConfig(device=None, strategy='auto', num_workers=None, pin_memory=None, ddp_timeout_seconds=None) dataclass

Device / dataloader runtime hints.

Note

- device=None + strategy="auto": choose "cuda" if available, else "cpu" with warning.
- strategy="auto" matches trainer logic (DDP if >1 GPU).
- ddp_timeout_seconds controls the distributed process-group timeout for
  DDP collectives. If None, DDPStrategy reads DVT_DDP_TIMEOUT_SECONDS
  when present, otherwise it uses its internal default.

OptimConfig(optimizer=None, optim_args=dict(), param_groups='flat', weight_decay=None, layer_decay=None) dataclass

Optimizer configuration. This configuration allows to generate torch.optim objects from yaml and/or export them to yaml. Handles LLRD strategies.

???+ note "Note": Allowed values for optimizer: - ImportSpec (portable) - torch optimizer class/type (portable if importable, e.g. torch.optim.AdamW) Disallowed: - optimizer instances (runtime-only, NOT portable) - lambdas / closures / non-importable callables (NOT portable) LLRD : - learning rate layer decay is an excellent optimization strategy when dealing with large backbones. It assigns decreasing learning rate to all blocks in the model.

Parameters:

Name	Type	Description	Default
optimizer	`Optional[Any]`	use optim type (python obj like torch.optim.Adamw) or ImportSpec.	None
optim_args	`Dict[str, Any]`	all kwargs style dict you want to feed the optimizer with.	dict()
param_groups	`str | Callable`	"flat" or "llrd" or user provided callable that build param groups from model. llrd will call deepvisiontools internal llrd strategy to compute param groups.	'flat'
weight_decay	`Optional[float]`	optional weight decay (regularization) for optimizer	None
layer_decay	Optional[float]	parameter used in llrd strategy for internal llrd strategy.	None

get_optimizer_spec()

Return an ImportSpec (portable), coercing from a class/type when needed. Raises if optimizer is missing or not portable.

Source code in src/deepvisiontools/train/trainconfig.py

def get_optimizer_spec(self) -> ImportSpec:
    """
    Return an ImportSpec (portable), coercing from a class/type when needed.
    Raises if optimizer is missing or not portable.
    """
    if self.optimizer is None:
        raise ValueError(
            "OptimConfig.optimizer must be set (ImportSpec or optimizer class)."
        )

    # Case 1: already portable
    if isinstance(self.optimizer, ImportSpec):
        return self.optimizer

    # Reject instances explicitly (not portable)
    try:
        import torch

        if isinstance(self.optimizer, torch.optim.Optimizer):
            raise TypeError(
                "OptimConfig.optimizer must be an optimizer CLASS (e.g. torch.optim.AdamW) "
                "or an ImportSpec. Optimizer INSTANCES are not portable and are rejected."
            )
    except Exception:
        # If torch isn't available, we still enforce class/type below.
        pass

    # Case 2: optimizer class/type (preferred ergonomic API)
    # We require an actual type to avoid lambdas/closures that can't be imported.
    if isinstance(self.optimizer, type):
        target = self._callable_to_target(self.optimizer)
        if target is None:
            raise TypeError(
                "OptimConfig.optimizer is a type but not importable (missing __module__/__qualname__). "
                "Use ImportSpec(target=..., kwargs=...)."
            )

        # Coerce and store back so export/describe become stable
        spec = ImportSpec(target=target, kwargs={})
        self.optimizer = spec
        return spec

    # Anything else is rejected (non-portable)
    raise TypeError(
        "OptimConfig.optimizer must be either an ImportSpec or an optimizer CLASS (type). "
        f"Got: {type(self.optimizer)}"
    )

SchedConfig(scheduler=None, scheduler_args=dict(), step_on='epoch', warmup=None) dataclass

Optim scheduler configuration

Parameters:

Name	Type	Description	Default
scheduler	`Optional[ImportSpec]`	siùmilar to optimizer, needs to be an importable object. Defaults to None (lr won't change).	None
scheduler_args	`Optional[Dict[str, Any]]`	args used in scheduler. Defaults to None.	dict()
step_on	`Literal["epoch", "step", "batch"]`	when lr is updated by scheduler. Defaults to epoch	'epoch'
warmup	`Optional[WarmupConfig]`	if warmup is required. Defaults to None.	None

get_scheduler_spec()

Return an ImportSpec (portable), coercing from a class/type when needed.

Raises if the scheduler is not portable or was assigned to the wrong field.

Source code in src/deepvisiontools/train/trainconfig.py

def get_scheduler_spec(self) -> ImportSpec | None:
    """
    Return an ImportSpec (portable), coercing from a class/type when needed.

    Raises if the scheduler is not portable or was assigned to the wrong field.
    """
    if self.scheduler is None:
        return None

    # Case 1: already portable
    if isinstance(self.scheduler, ImportSpec):
        return self.scheduler

    # Reject runtime instances explicitly
    if self._is_scheduler_instance(self.scheduler):
        raise TypeError(
            "SchedConfig.scheduler must be a scheduler CLASS "
            "(e.g. torch.optim.lr_scheduler.ExponentialLR) or an ImportSpec. "
            "Scheduler INSTANCES are runtime-only and are rejected."
        )

    # Case 2: scheduler class/type (preferred ergonomic API)
    if isinstance(self.scheduler, type):
        target = self._callable_to_target(self.scheduler)
        if target is None:
            raise TypeError(
                "SchedConfig.scheduler is a type but not importable "
                "(missing __module__/__qualname__). "
                "Use ImportSpec(target=..., kwargs=...)."
            )

        # Coerce and store back so export/describe become stable
        spec = ImportSpec(target=target, kwargs={})
        self.scheduler = spec
        return spec

    # Common user mistake: args dict assigned to `scheduler` instead of `scheduler_args`
    if isinstance(self.scheduler, dict):
        raise TypeError(
            "SchedConfig.scheduler received a dict. "
            "The scheduler itself must be an ImportSpec or scheduler CLASS; "
            "its kwargs belong in `sched.scheduler_args`."
        )

    # Anything else is rejected
    raise TypeError(
        "SchedConfig.scheduler must be either None, an ImportSpec, or a scheduler CLASS (type). "
        f"Got: {type(self.scheduler)}"
    )

build_scheduler(optimizer)

Validate/coerce the scheduler and instantiate it.

Calling get_scheduler_spec() first ensures: - class/type inputs are converted to ImportSpec for stable export - invalid assignments fail early with a helpful message

Source code in src/deepvisiontools/train/trainconfig.py

def build_scheduler(self, optimizer: Any) -> Any:
    """
    Validate/coerce the scheduler and instantiate it.

    Calling `get_scheduler_spec()` first ensures:
      - class/type inputs are converted to ImportSpec for stable export
      - invalid assignments fail early with a helpful message
    """
    spec = self.get_scheduler_spec()
    if spec is None:
        return None

    # Use the now-coerced portable representation for instantiation
    return self(optimizer)

WarmupConfig(kind='linear', steps=None, epochs=None, steps_unit='step', ratio=1.0, start_factor=0.001) dataclass

Warmup configuration.

Parameters:

Name	Type	Description	Default
kind	`Literal["linear", "constant"]`	Default to "linear".	'linear'
steps	`Optional[int]`	Default to None	None
epochs	`Optional[float]`	Default to None	None
steps_unit	`Literal["step", "epoch"]`	Default to"step"	'step'
ratio	`float`	linear warmup starts from start_factor * (ratio * base_lr). Defaults to 1.0	1.0
start_factor	`float`	Defaults to 0.001	0.001

???+ note Note: Supports 2 ways to specify duration: - steps: number of optimizer steps (batch-based warmup) - epochs: number of epochs (converted to steps by Trainer using len(train_loader)) A "switch" is provided via steps_unit: - steps_unit="step" -> interpret steps as optimizer steps (default) - steps_unit="epoch" -> interpret steps as epochs (converted into epochs) so WarmupConfig(steps=1, steps_unit="epoch") == WarmupConfig(epochs=1) - Trainer will compute warmup_steps as: - if steps is not None: warmup_steps = steps - elif epochs is not None: warmup_steps = epochs * len(train_loader) - So do NOT set both steps and epochs (we normalize to one).

PrecisionConfig(amp=None, amp_dtype='bf16', grad_clip=None, accumulate_steps=1) dataclass

Torch AMP and gradient precision configuration.

Automatic mixed precision is strongly recommended when supported: it is usually faster, uses less CUDA memory, and often has no accuracy trade-off. This config also controls gradient clipping and gradient accumulation.

Parameters:

Name	Type	Description	Default
amp	`bool | None`	Enable or disable AMP. If None, the Trainer falls back to the global current() config. Defaults to None.	None
amp_dtype	`Literal["fp16", "bf16", "fp32"]`	Precision dtype used by AMP. Defaults to "bf16".	'bf16'
grad_clip	`float | None`	Optional gradient clipping value. Typical values are in the 1-10 range for very large models. Defaults to None.	None
accumulate_steps	`int`	Number of batches to accumulate before optimizer update, scheduler step, and EMA update. Useful to emulate a larger effective batch size when limited by hardware memory. Defaults to 1.	1

EMAConfig(enabled=True, decay=0.995, update_every=1, warmup_steps=0, eval_with_ema=True) dataclass

Exponential moving average configuration.

EMA keeps a smoothed copy of model weights during training. After each optimizer step, the EMA state is updated from the current model weights. EMA weights are often more stable for validation and deployment, especially for noisy training setups such as detection and instance segmentation.

Parameters:

Name	Type	Description	Default
enabled	`bool`	Whether to enable EMA tracking. Defaults to True.	True
decay	`float`	EMA decay factor. Typical values range from 0.99 for short trainings to 0.9998 for long trainings. Defaults to 0.995.	0.995
update_every	`int`	Update EMA every N optimizer steps. Defaults to 1.	1
warmup_steps	`int`	Number of optimizer steps used to ramp EMA decay at the beginning of training. Defaults to 0.	0
eval_with_ema	`bool`	Whether to temporarily swap EMA weights during validation/evaluation. Defaults to True.	True

FreezingConfig(freeze_epochs=0, backbone_group_name='backbone', freeze_backbone=False, freeze_bn=False) dataclass

If you need to freeze backbone at begining of training

NormalizationConfig(policy='auto', freeze_affine=True, freeze_running_stats=True, verbose=True) dataclass

BatchNorm / SyncBatchNorm policy. Used in DDP and small batches to freeze batch norm layers / sync BN.

Policies

• "auto": Safe default.

  • if no BN is present -> do nothing
  • if BN is present and local micro-batch <= 1 -> freeze BN
  • otherwise -> do nothing This intentionally does NOT auto-convert to SyncBN.

• "none": Never touch normalization layers.

• "freeze_bn": Freeze BatchNorm / SyncBatchNorm affine params (optional) and keep running stats frozen during train epochs.

• "sync_bn": Convert BatchNorm{1,2,3}d to SyncBatchNorm, but only when DDP is active on CUDA. Otherwise the request is ignored with a logged summary.

• "sync_or_freeze_bn": If DDP+CUDA is active -> convert to SyncBN. Otherwise -> freeze BN.

Notes

• LayerNorm / GroupNorm / InstanceNorm are never modified by this policy.
• "freeze BN" is implemented by:
  • setting affine params requires_grad=False (if freeze_affine=True)
  • forcing BN modules to eval() whenever the trainer enters train mode (if freeze_running_stats=True)

CheckpointConfig(output_dir=None, enable_checkpointing=True, monitor=('valid/loss', 'loss'), mode='min', filename='checkpoint', save_every_n_epochs=1, keep_last_n=3, export_model_artifact=False, export_when='best', model_builder=None, model_kwargs=None) dataclass

Handles checkpointing and optional model artifact export.

Checkpointing is used to save restartable training states and to keep track of the best model according to a monitored metric.

Parameters:

Name	Type	Description	Default
output_dir	`str | None`	Optional checkpoint output directory. If None, the Trainer uses the default deepvisiontools output setup.	None
enable_checkpointing	`bool`	Master switch for checkpointing and model artifact export. If False, training and logging still run normally, but checkpoint/model export phases are skipped.	True
monitor	`tuple[str, str]`	Metric to monitor as (metric_complete_name, metric_type). Defaults to ("valid/loss", "loss").	('valid/loss', 'loss')
mode	`Literal["min", "max"]`	Direction used to decide whether a model is better. Use "min" for losses and "max" for metrics to maximize.	'min'
filename	`str`	Prefix used for checkpoint filenames. Defaults to "checkpoint".	'checkpoint'
save_every_n_epochs	`int`	Periodic checkpoint interval in epochs. If 0, no periodic checkpoint is written. Defaults to 1.	1
keep_last_n	`int`	Number of most recent periodic checkpoints to keep. Older periodic checkpoints are removed. Defaults to 3.	3
export_model_artifact	`bool`	Whether to export a portable model artifact in addition to training checkpoints. Defaults to False.	False
export_when	`Literal["best", "all"]`	Controls when model artifacts are exported: only for the best model or at every export opportunity.	'best'
model_builder	`str | None`	Optional model builder key saved as a restore hint. Usually inferred automatically from model_spec.	None
model_kwargs	`dict[str, Any] | None`	Optional model builder kwargs saved as restore hints. Usually inferred automatically from model_spec.	None

Note

• enable_checkpointing=False is the master switch. When disabled, the training loop still runs normally and MLflow/TensorBoard logging still works, but checkpoint/model export phases are skipped.
• This is especially useful for large Optuna studies where you only want metrics/losses/logging without writing restart artifacts.

LoggingConfig(backends='mlflow', tensorboard=dict(), mlflow=dict(), log_every_n_steps=0, mlflow_experiment_name=None, mlflow_run_name=None, wandb_project=None, wandb_run_name=None) dataclass

Monitoring configuration.

MLflow can already consume arbitrary keys from mlflow, but exposing the most common ones explicitly makes the Trainer config much cleaner and safer.

Parameters:

Name	Type	Description	Default
backends	str | collections.abc.Sequence[str]	Logging backend(s), e.g. "mlflow", "tensorboard", or both.	'mlflow'
tensorboard	dict[str, typing.Any]	Extra kwargs forwarded to TensorBoard logger.	dict()
mlflow	dict[str, typing.Any]	Extra kwargs forwarded to MLflowLogger.	dict()
log_every_n_steps	int	Logging period hint used by hooks.	0
mlflow_experiment_name	str | None	Optional explicit MLflow experiment name.	None
mlflow_run_name	str | None	Optional explicit MLflow run name.	None
wandb_project	str | None	Reserved/legacy field.	None
wandb_run_name	str | None	Reserved/legacy field.	None

EvalConfig(eval_every_n_steps=0, eval_on='valid', metrics=None) dataclass

Evaluation configuration.

Parameters:

Name	Type	Description	Default
eval_every_n_steps	`int`	Optional step-based evaluation period. If 0, evaluation is performed once per epoch. Defaults to 0.	0
eval_on	`Literal["valid", "train", "valid_and_train"]`	Dataset split(s) used for evaluation. Evaluating on train during training is possible but slower. If "valid_and_train" passed, internally clone trainloader, switch off augmentations, and evaluate on it. Defaults to "valid".	'valid'
metrics	`MetricsSpec | None`	Portable metric suite configuration. If None, no metrics factory is built. Defaults to None.	None

MetricsSpec(kind=None, num_classes=None, target=None, kwargs=dict()) dataclass

Main metric setup config class used in EvalConfig. Two modes: 1) kind-based: kind in {"bbox","instance_mask","semantic_mask"} + num_classes 2) target-based: import target "module:qualname" + kwargs (if having special callables). Note that this second mode had not been widely tested yet.

Parameters:

Name	Type	Description	Default
kind	`Optional[str]`	str val from "detection", "semantic", "instance". Setup standard metric suites in deepvisiontools	None
num_classes	`Optional[int]`	if you want to override, but dvt computes class from current() when needed. Defaults to None	None
target	Optional[str]	for specifi metrics callables.	None
kwargs	Optional[Dict[str, typing.Any]]	args for your specific metric	dict()

build_factory()

Return a zero-arg factory that builds the metric suite object.

Rules: - If target is set: import + kwargs - Else: kind-based factory, with num_classes taken from: - self.num_classes if provided - else current().data_* defaults depending on kind

Source code in src/deepvisiontools/train/trainconfig.py

def build_factory(self) -> Callable[[], Any]:
    """Return a zero-arg factory that builds the metric suite object.

    Rules:
    - If target is set: import + kwargs
    - Else: kind-based factory, with num_classes taken from:
        - self.num_classes if provided
        - else current().data_* defaults depending on kind
    """
    # ---- explicit import target wins ----
    if self.target:
        obj = _import_from_target(self.target)

        def _f():
            return obj(**dict(self.kwargs))

        return _f

    # ---- kind-based default mapping ----
    kind = (self.kind or "").strip().lower()
    if not kind:

        def _none():
            raise RuntimeError(
                "MetricsSpec is empty (no target and no kind). "
                "Set eval.metrics to null, or provide target or kind."
            )

        return _none

    # Resolve num_classes:
    # - prefer explicit num_classes
    # - else fallback to global config defaults
    nc_val = self.num_classes
    if nc_val is None:
        cfg = (
            current()
        )  # from ..config import current (already imported in trainconfig.py)
        if kind in {"bbox", "detection", "det"}:
            nc_val = int(cfg.data_box.num_classes)
        elif kind in {"instance_mask", "instance", "inst", "mask"}:
            nc_val = int(cfg.data_instance_mask.num_classes)
        elif kind in {"semantic_mask", "semantic", "semseg"}:
            nc_val = int(cfg.data_semantic_mask.num_classes)
        else:
            raise ValueError(
                f"Unknown MetricsSpec.kind={kind!r}. Expected one of: "
                f"'bbox', 'instance_mask', 'semantic_mask' (or aliases). "
                f"Alternatively provide MetricsSpec(target=..., kwargs=...)."
            )

    nc = int(nc_val)
    if nc <= 0:
        raise ValueError(f"num_classes must be > 0, got {nc}")

    # Local import to keep config module lightweight
    from ..metrics.factories import (
        DetectionSuiteFactory,
        InstanceSegSuiteFactory,
        SemanticSegSuiteFactory,
    )

    # Accept a few aliases
    if kind in {"bbox", "detection", "det"}:

        def _f():
            return DetectionSuiteFactory(nc)()

        return _f

    if kind in {"instance_mask", "instance", "inst", "mask"}:

        def _f():
            return InstanceSegSuiteFactory(nc)()

        return _f

    if kind in {"semantic_mask", "semantic", "semseg"}:
        keep_per_image = bool(self.kwargs.get("keep_per_image", False))

        def _f():
            return SemanticSegSuiteFactory(nc, keep_per_image=keep_per_image)()

        return _f

    raise ValueError(
        f"Unknown MetricsSpec.kind={kind!r}. Expected one of: "
        f"'bbox', 'instance_mask', 'semantic_mask' (or aliases). "
        f"Alternatively provide MetricsSpec(target=..., kwargs=...)."
    )