class TrainConfig(BaseConfig):
"""Training hyperparameters and auto-batching configuration.
Notes:
* ``auto_batch_target_effective`` is interpreted as the **global**
effective batch size target. In multi-GPU / multi-node runs the auto-batch resolver derives a per-device
target by dividing this value across ``devices * num_nodes`` before selecting ``grad_accum_steps``.
"""
# extra="forbid" arms BaseConfig.catch_typo_kwargs so typo'd train() kwargs (e.g. ``epoch`` instead of
# ``epochs``) raise with a helpful message instead of being silently ignored. Legacy kwargs handled by
# RFDETR.train() (resolution/device/callbacks/start_epoch/do_benchmark) are popped before construction.
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", validate_assignment=True)
lr: float = 1e-4
lr_encoder: float = 1.5e-4
batch_size: int | Literal["auto"] = 4
# Gradient accumulation is an explicit opt-in, not a silent default: max out batch_size for the GPU first
# and raise this only when memory forces a smaller physical batch. Defaulting to 1 (was 4) drops the
# default effective batch from 16 to 4 — a training-semantics change, see CHANGELOG. Runs with
# batch_size="auto" are unaffected: the auto-batch probe overwrites this field (see RFDETR.train).
grad_accum_steps: int = 1
# Batch size for the validation, test and predict dataloaders. None (the default) inherits the resolved
# train batch size, i.e. exactly the pre-existing behavior. Evaluation runs under no_grad, so it avoids
# autograd activation storage, but in-fit validation still shares device memory with the model and optimizer
# state and needs memory for its own forward outputs. A train batch_size lowered to fit an optimizer step can
# therefore still unnecessarily shrink eval batches. The `eval_` prefix (not `val_`) matches the other
# evaluation-side knobs here (eval_ema_only, eval_max_dets, eval_interval, eval_masks_head_resolution)
# and reflects that this governs all three eval loaders, not validation alone. Unlike batch_size it
# accepts no "auto": it is never probed, and an explicit value stays usable even when batch_size="auto"
# has not been resolved.
eval_batch_size: int | None = None
# Global effective batch size target, divided across devices and nodes. This is a floor, not a cap: the probe
# only raises grad_accum_steps to *reach* it (see recommend_grad_accum_steps), and never shrinks the micro-batch
# to hold it. Once the probed micro-batch already meets or exceeds this value, grad_accum_steps stays at 1 and
# the effective batch is simply whatever fit in memory — so it grows with VRAM while lr stays put. Pin
# batch_size to a concrete integer when training semantics must match across GPUs of different sizes.
auto_batch_target_effective: int = 16
# Auto-batch probe: worst-case assumptions when batch_size="auto".
auto_batch_max_targets_per_image: int = 100
auto_batch_ema_headroom: float = 0.7 # scale safe batch by this when use_ema=True (EMA uses extra memory)
epochs: int = 100
resume: PathLikeStr | None = None
ema_decay: float = 0.993
ema_tau: int = 100
lr_drop: int = 100
checkpoint_interval: int = Field(default=10, ge=1)
skip_best_epochs: int = Field(default=0, ge=0)
smooth_alpha: float = 0.0
warmup_epochs: float = 0.0
lr_vit_layer_decay: float = 0.8
lr_component_decay: float = 0.7
drop_path: float = 0.0
cls_loss_coef: float = 1.0
# Detection-vs-keypoint distinction is derived by callers via `include_keypoints`, not
# stored on this field. See rfdetr.datasets.transforms.AlbumentationsWrapper.from_config
# for the None/[]/[...] tri-state contract applied at the augmentation-pipeline boundary.
keypoint_flip_pairs: list[int] = Field(default_factory=list)
keypoint_l1_loss_coef: float = 0
keypoint_findable_loss_coef: float = 0
keypoint_visible_loss_coef: float = 0
keypoint_nll_loss_coef: float = 0
keypoint_oks_sigmas: list[float] | None = None
dataset_file: Literal["coco", "o365", "roboflow", "yolo"] = "roboflow"
square_resize_div_64: bool = True
dataset_dir: PathLikeStr | None
output_dir: PathLikeStr = "output"
# XLA/TPU: every distinct (H, W) triggers a separate graph compilation. Set multi_scale=False
# for a static shape (zero recompilations after the first batch) when training on TPU.
multi_scale: bool = True
expanded_scales: bool = True
do_random_resize_via_padding: bool = False
use_ema: bool = True
ema_update_interval: int = 1
# Validation-only: also evaluate the base model, on top of the model validation already forwards
# through. Validation evaluates exactly ONE model by default — the EMA-averaged weights when
# use_ema=True, the base weights otherwise — because the EMA model is the one best-checkpoint
# selection ships and the base-model pass is diagnostic. Dropping it removes one full forward pass
# over the validation set per epoch (measured ~3-3.5% of epoch time). Set eval_base_model=True to
# restore the base+EMA comparison: validation_step then forwards the base model and COCOEvalCallback
# runs the EMA model in a second no_grad pass, exactly as it did before this default existed.
# Inert when use_ema=False — the base model is the selected model, so it is evaluated either way.
# Metric-key routing: val/mAP_50_95 (and val/mAP_50, val/mAR, val/segm_mAP_*) always report the
# model that was evaluated first-class — the base model when eval_base_model=True, otherwise the
# selected (EMA) model, mirrored from the EMA track by COCOEvalCallback so every scheduler,
# early-stopping hook, checkpoint monitor and dashboard watching the primary key keeps receiving a
# real number. val/ema_* is always the EMA track and is what the EMA checkpoint monitor reads.
# The EMA-only path mirrors aggregate mAP/mAR and per-class AP onto the primary namespace; val/ema_* remains
# available for explicit EMA monitors. The printed summary table is titled "val (ema)" when the base model was
# not evaluated. val/F1 has no parallel EMA-tracked
# accumulator and always follows validation_step's own forward — under the default that is the EMA
# model, which now agrees with the mAP under the same key.
eval_base_model: bool = False
# Deprecated: superseded by eval_base_model (removal in v1.13). Evaluating only the EMA model is now
# the default, so this no-op flag remains through the 0.3-cycle deprecation window; setting it emits a
# FutureWarning. It still requires use_ema=True and contradicts eval_base_model=True.
eval_ema_only: bool = False
num_workers: int = 2
weight_decay: float = 1e-4
amp_dtype: Literal["auto", "bf16", "fp16"] = Field(
default="auto",
description=(
"Mixed-precision autocast dtype. "
"'auto' selects bf16-mixed on Ampere+ CUDA, fp16 otherwise. "
"'bf16' forces bfloat16 (falls back to fp16 with a warning if unsupported). "
"'fp16' forces fp16. "
"Has no effect when model_config.amp=False or when training on CPU."
),
)
best_model_metric: Literal["map", "mar"] = Field(
default="map",
description=(
"Validation metric that selects the best checkpoint, and (when early_stopping=True) "
"that early stopping monitors. 'map' (default) uses the task's mAP@50:95 (keypoint OKS "
"AP, segmentation mask AP, or box AP). 'mar' uses box mAR at the configured eval_max_dets "
"for detection and segmentation, and keypoint mAR (OKS-based) at fixed COCO maxDets=20. "
"Segmentation mAR is box-level because torchmetrics does not expose a separate mask mAR."
),
)
early_stopping: bool = False
early_stopping_patience: int = 10
early_stopping_min_delta: float = 0.001
early_stopping_use_ema: bool = False
progress_bar: Literal["tqdm", "rich"] | None = None # Progress bar style: "rich", "tqdm", or None to disable.
tensorboard: bool = True
wandb: bool = False
mlflow: bool = False
clearml: bool = False # Not yet implemented — reserved for future use.
project: str | None = None
run: str | None = None
class_names: list[str] | None = None
run_test: bool = False
eval_max_dets: int = 500
eval_interval: int = 1
log_per_class_metrics: bool = False
# Segmentation only. Skip upsampling predicted masks to full image resolution during
# validation/test, returning them at the mask head's native (lower) resolution instead —
# cheaper, but ground-truth masks must then be compared at that same lower resolution
# (handled in COCOEvalCallback). No effect on non-segmentation models or on inference output
# (RFDETR.predict always upsamples regardless of this flag).
# Metric comparability: GT masks are nearest-downsized to the mask head's native resolution
# (e.g. 512x512 -> ~16x16) before comparison, so val/segm_mAP under this flag is NOT
# comparable to a full-resolution run — small objects can collapse to empty masks at the
# lower resolution, and IoU is computed on a coarser pixel grid either way.
eval_masks_head_resolution: bool = False
aug_config: Optional[Dict[str, Any]] = None
scale_jitter: bool = True
augmentation_backend: AugmentationBackend | Literal["cpu", "auto"] = "cpu"
save_dataset_grids: bool = False
@field_validator("augmentation_backend", mode="before")
@classmethod
def _coerce_augmentation_backend(cls, v: Any) -> Any:
"""Map legacy backend name strings (``"gpu"``, ``"tv"``, ``"albu"``) to their current form.
``"cpu"``/``"auto"`` pass through unchanged — they are auto-pick sentinels resolved lazily by
:meth:`AugmentationBackend.from_str` at dataset-build time, not at config construction time, so a saved config
stays portable across environments. See :class:`AugmentationBackend` for the full rationale.
"""
if isinstance(v, str):
return _LEGACY_AUGMENTATION_BACKEND_ALIASES.get(v, v)
return v
@field_serializer("augmentation_backend")
def _serialize_augmentation_backend(self, value: AugmentationBackend | str) -> str:
"""Serialize the backend selector to its plain string form.
Returns the enum's ``.value`` for :class:`AugmentationBackend` members and passes the
``"cpu"``/``"auto"`` sentinel strings through unchanged. This lets checkpoint writers call
plain :meth:`model_dump` and still get a JSON-safe ``str`` for this one field, instead of a
blanket ``model_dump(mode="json")`` that would also coerce every *other* field's serialized
shape (e.g. the ``int`` keypoint loss coefficients to ``float``). Applies in both python and
json ``model_dump`` modes, so the two checkpoint writers stay consistent.
Args:
value: Stored backend selector — an :class:`AugmentationBackend` member or a
``"cpu"``/``"auto"`` sentinel string.
Returns:
The plain string form of the backend selector.
Examples:
>>> cfg = TrainConfig(dataset_dir="ds", augmentation_backend="torchvision")
>>> cfg.model_dump()["augmentation_backend"]
'torchvision'
"""
if isinstance(value, AugmentationBackend):
return str(value.value)
return value
notes: Optional[Any] = Field(
default=None,
description=(
"User-defined provenance metadata embedded in best-model .pth checkpoints "
"under checkpoint['args']['notes'] and in exported ONNX files under the "
"'rfdetr_notes' metadata property. Accepts any JSON-serialisable value "
"(string, dict, list, int, float, bool). String values are stored verbatim; "
"all other types are JSON-encoded."
),
)
@field_validator("progress_bar", mode="before")
@classmethod
def _coerce_legacy_progress_bar(cls, value: Any) -> Any:
"""Normalize legacy boolean progress_bar values to the new string/None representation.
This preserves compatibility with older configs where ``progress_bar`` was a bool.
"""
if isinstance(value, bool):
return "tqdm" if value else None
return value
@field_validator("amp_dtype", mode="before")
@classmethod
def _coerce_amp_dtype(cls, value: Any) -> Any:
"""Fall back to ``'auto'`` (with a warning) for an unrecognised or wrong-typed ``amp_dtype``.
Mixed precision is a best-effort speed/memory optimisation, so an invalid request degrades to the auto-selected
dtype rather than failing the whole training run.
"""
if value not in ("auto", "bf16", "fp16"):
# stacklevel=2 points into Pydantic internals; unavoidable with @field_validator in Pydantic v2.
warnings.warn(
f"Unknown amp_dtype={value!r}; expected one of 'auto', 'bf16', 'fp16'. Falling back to 'auto'.",
UserWarning,
stacklevel=2,
)
return "auto"
return value
# Promoted from populate_args() — PTL migration (T4-2).
# device is intentionally absent: PTL auto-detects accelerator via Trainer(accelerator="auto").
accelerator: str = "auto"
clip_max_norm: float = 0.1
seed: int | None = None
sync_bn: bool = False
# strategy maps to PTL Trainer(strategy=...). Common values: "auto", "ddp",
# "ddp_spawn", "fsdp", "deepspeed". Invalid values surface as PTL errors.
strategy: str = "auto"
devices: int | str = 1
# num_nodes maps to PTL Trainer(num_nodes=...) for multi-machine training.
# Single-machine DDP users should leave this at 1 (the default).
num_nodes: int = 1
fp16_eval: bool = False
lr_scheduler: str | Callable[..., SchedulerType] = "step"
lr_scheduler_kwargs: dict[str, Any] = Field(default_factory=dict)
lr_scheduler_interval: Literal["step", "epoch"] = "step"
lr_scheduler_monitor: str = "val/loss"
# Deprecated aux LR knobs — kept for one cycle; folded into lr_scheduler_kwargs (see _map_deprecated_lr_fields).
lr_min_factor: float = 0.0
optimizer: str | Callable[..., Optimizer] = "adamw"
optimizer_kwargs: dict[str, Any] = Field(default_factory=dict)
dont_save_weights: bool = False
# PTL runtime/perf tuning knobs.
train_log_sync_dist: bool = False
# Component-level train/ metrics honor train_log_on_step for on_step visibility only when
# compact_train_metrics=False; when True, components are aggregated and logged on_epoch-only regardless.
train_log_on_step: bool = False
# When True (default), per-decoder/encoder auxiliary loss keys (e.g. loss_ce_0, loss_bbox_enc) are
# aggregated into train/<term>_aux and always logged on_epoch-only. When False, every per-layer key is
# logged individually, honoring train_log_on_step for those component logs too.
compact_train_metrics: bool = True
compute_train_metrics: bool = False
# Restores PTL's pre-training sanity-validation pass (0 = disabled, current default;
# increase to re-enable and catch val-path errors before a full epoch runs).
num_sanity_val_steps: int = 0
compute_val_loss: bool | Literal["auto"] = "auto"
# No "auto" here: unlike val/loss, nothing (schedulers, callbacks) monitors test/loss, and
# trainer.test() runs once rather than every epoch, so consumer-based auto-detection doesn't apply.
compute_test_loss: bool = True
pin_memory: bool | None = None
persistent_workers: bool | None = None
prefetch_factor: int | None = None
pack_targets: bool = True
@field_validator("batch_size", mode="after")
@classmethod
def validate_batch_size(cls, v: int | Literal["auto"]) -> int | Literal["auto"]:
"""Validate batch_size is a positive integer or the literal 'auto'."""
if v == "auto":
return v
if v < 1:
raise ValueError("batch_size must be >= 1, or 'auto'.")
return v
@field_validator("eval_batch_size", mode="after")
@classmethod
def validate_eval_batch_size(cls, v: int | None) -> int | None:
"""Validate eval_batch_size is None (inherit the train batch size) or >= 1."""
if v is not None and v < 1:
raise ValueError("eval_batch_size must be >= 1 when provided.")
return v
@field_validator(
"grad_accum_steps", "auto_batch_target_effective", "auto_batch_max_targets_per_image", mode="after"
)
@classmethod
def validate_positive_train_steps(cls, v: int) -> int:
"""Validate accumulation, target-effective batch, and max targets are >= 1."""
if v < 1:
raise ValueError(
"grad_accum_steps, auto_batch_target_effective, and auto_batch_max_targets_per_image must be >= 1."
)
return v
@field_validator("auto_batch_ema_headroom", mode="after")
@classmethod
def validate_ema_headroom(cls, v: float) -> float:
"""Validate auto_batch_ema_headroom is in (0, 1]."""
if not (0 < v <= 1.0):
raise ValueError("auto_batch_ema_headroom must be in (0, 1].")
return v
@field_validator("smooth_alpha", mode="after")
@classmethod
def validate_smooth_alpha(cls, v: float) -> float:
"""Validate smooth_alpha is in [0.0, 1.0)."""
if not (0.0 <= v < 1.0):
raise ValueError("smooth_alpha must be in [0.0, 1.0).")
return v
@field_validator("ema_update_interval", "eval_interval", mode="after")
@classmethod
def validate_positive_intervals(cls, v: int) -> int:
"""Validate interval fields are >= 1."""
if v < 1:
raise ValueError("Interval fields must be >= 1.")
return v
@model_validator(mode="before")
@classmethod
def _desugar_callable_optimizer(cls, data: Any) -> Any:
"""Desugar a reconstructable callable optimizer into its serializable string form.
A callable ``optimizer`` (a class or ``functools.partial``) that can be imported is rewritten to a dotted import
path plus ``optimizer_kwargs`` so the config round-trips through ``training_config.json``. User-supplied
``optimizer_kwargs`` are ignored for callable optimizers (bake arguments into the callable instead). Non-
reconstructable callables are kept as-is and only warned about.
"""
if not isinstance(data, dict):
return data
optimizer = data.get("optimizer")
if optimizer is None or isinstance(optimizer, str) or not callable(optimizer):
return data
if data.get("optimizer_kwargs"):
warnings.warn(
"optimizer_kwargs is ignored when optimizer is a callable; bake arguments into the "
"callable (for example with functools.partial) instead.",
stacklevel=2,
)
path, kwargs, reason = _desugar_optimizer_callable(optimizer)
if reason is None:
data["optimizer"] = path
data["optimizer_kwargs"] = kwargs
else:
data["optimizer_kwargs"] = {}
label = getattr(optimizer, "__qualname__", None) or repr(optimizer)
warnings.warn(
f"optimizer callable {label!r} cannot be saved to training_config.json and restored: "
f"{reason}. Training proceeds with the in-memory callable; only saved-config "
"reproducibility is affected.",
stacklevel=2,
)
return data
@field_validator("optimizer", mode="after")
@classmethod
def validate_optimizer_name(cls, v: str | Callable[..., Optimizer]) -> str | Callable[..., Optimizer]:
"""Validate a string optimizer: a bare name must be a native torch.optim optimizer."""
if not isinstance(v, str):
return v
optimizer = v.strip()
if not optimizer:
raise ValueError("optimizer must be a non-empty string.")
# Bare short names must resolve to a torch.optim optimizer (checked eagerly).
# Dotted import paths are validated lazily at train start (the module may be optional).
if "." not in optimizer:
_resolve_native_optimizer(optimizer)
return optimizer
@model_validator(mode="after")
def validate_eval_ema_only(self) -> "TrainConfig":
"""``eval_ema_only`` has no EMA model to evaluate without ``use_ema=True``, and contradicts ``eval_base_model``.
The flag is deprecated (see ``_warn_deprecated_eval_ema_only``) but still validated: absorbing a contradictory
pair into the no-op alias would leave one of the two settings silently without effect.
"""
if self.eval_ema_only and not self.use_ema:
raise ValueError("eval_ema_only=True requires use_ema=True.")
if self.eval_ema_only and self.eval_base_model:
raise ValueError(
"eval_ema_only=True contradicts eval_base_model=True. Drop the deprecated eval_ema_only: "
"evaluating only the selected model is now the default."
)
return self
@model_validator(mode="before")
@classmethod
def _warn_deprecated_eval_ema_only(cls, data: Any) -> Any:
"""Warn that ``eval_ema_only`` is superseded by the default single-model evaluation policy.
Legacy input that contains ``eval_ema_only`` without the new ``eval_base_model`` field is migrated to the
equivalent old base-plus-EMA policy and warns. New dumps contain both fields, so reloading them stays silent.
"""
if not isinstance(data, dict):
return data
if "eval_ema_only" not in data:
return data
if "eval_base_model" in data and not data["eval_base_model"]:
return data
data = dict(data)
if "eval_base_model" not in data:
data["eval_base_model"] = not bool(data["eval_ema_only"])
if data.get("eval_ema_only") is not None:
warnings.warn(
"eval_ema_only is deprecated. New configurations evaluate only the selected model "
"(EMA when use_ema=True) by default; legacy configurations are migrated from this flag. "
"Set eval_base_model=True to also evaluate the base model.",
FutureWarning,
stacklevel=2,
)
return data
@model_validator(mode="after")
def validate_explicit_val_loss_disable(self) -> "TrainConfig":
"""Reject disabling validation loss for a configured plateau loss monitor.
Reconstructable scheduler callables are normalized to their dotted path before this validator runs. Non-
reconstructable callables are checked after their concrete scheduler is created by ``RFDETRModelModule``.
"""
if (
self.compute_val_loss is False
and self.lr_scheduler == "torch.optim.lr_scheduler.ReduceLROnPlateau"
and self.lr_scheduler_monitor == "val/loss"
):
raise ValueError(
"compute_val_loss=False requires a non-val/loss monitor when "
"lr_scheduler is ReduceLROnPlateau. Set compute_val_loss=True or 'auto'."
)
return self
@model_validator(mode="after")
def validate_optimizer_kwargs(self) -> "TrainConfig":
"""Reserved optimizer kwargs are only rejected for managed (short-name) optimizers."""
if _is_managed_optimizer_name(self.optimizer):
reserved_present = _OPTIMIZER_MANAGED_KWARGS.intersection(self.optimizer_kwargs)
if reserved_present:
reserved = ", ".join(sorted(reserved_present))
raise ValueError(f"optimizer_kwargs cannot include RF-DETR-managed key(s): {reserved}.")
return self
@model_validator(mode="after")
def validate_lr_scheduler_kwargs(self) -> "TrainConfig":
"""Reject unknown ``lr_scheduler_kwargs`` keys for the managed ``"step"`` / ``"cosine"`` presets.
Managed presets consume only ``min_factor`` and ``lr_drop``; any other key would be silently ignored, so surface
it as an error (mirroring ``validate_optimizer_kwargs``). Explicit schedulers forward their kwargs verbatim to
the constructor and are left unchecked here.
"""
if _is_managed_scheduler_name(self.lr_scheduler):
unknown = set(self.lr_scheduler_kwargs) - _MANAGED_SCHEDULER_KWARGS
if unknown:
allowed = ", ".join(sorted(_MANAGED_SCHEDULER_KWARGS))
unknown_keys = ", ".join(sorted(unknown))
raise ValueError(
f"lr_scheduler_kwargs for a managed preset ({self.lr_scheduler!r}) accepts only "
f"{{{allowed}}}; unknown key(s): {unknown_keys}."
)
return self
@model_validator(mode="before")
@classmethod
def _map_deprecated_lr_fields(cls, data: Any) -> Any:
"""Fold the deprecated ``lr_drop`` / ``lr_min_factor`` fields into ``lr_scheduler_kwargs``.
These loose knobs are deprecated in favor of ``lr_scheduler_kwargs``. When either is supplied with a non-default
value for a managed preset (``"step"`` / ``"cosine"``), it is copied into ``lr_scheduler_kwargs`` (without
overriding an explicit kwarg) and a ``FutureWarning`` is emitted. Default values are ignored silently so round-
tripping a dumped config (which always carries these fields) never warns. For explicit (dotted-path / callable)
schedulers the deprecated fields are preset-specific and left untouched.
"""
if not isinstance(data, dict):
return data
# Only managed presets consume these knobs; default lr_scheduler ("step") is managed.
if not _is_managed_scheduler_name(data.get("lr_scheduler", "step")):
# Explicit / callable scheduler: these preset knobs are inert. Warn (never fold) when a non-default
# value is set so a stale lr_drop / lr_min_factor carried over from a managed config is not silently
# dropped — a reproducibility footgun when migrating a saved config to an explicit scheduler.
for field_name in _DEPRECATED_LR_FIELD_KWARGS:
if field_name in data and data[field_name] != cls.model_fields[field_name].default:
warnings.warn(
f"{field_name} is ignored for the explicit (non-managed) lr_scheduler "
f"{data.get('lr_scheduler')!r}; it only applies to the managed 'step'/'cosine' presets.",
FutureWarning,
stacklevel=2,
)
return data
kwargs = dict(data.get("lr_scheduler_kwargs") or {})
for field_name, kwarg_name in _DEPRECATED_LR_FIELD_KWARGS.items():
if field_name not in data:
continue
# A default value (common when reloading a dumped config) is a no-op: the managed builder falls back to
# the same default. Skip silently so serialization round-trips don't emit spurious deprecation warnings.
if data[field_name] == cls.model_fields[field_name].default:
continue
# Already migrated: a dumped config carries both the top-level field and the folded kwarg. If the kwarg
# already holds this value, the field adds nothing — skip silently so migrated-config reloads never warn.
if kwargs.get(kwarg_name) == data[field_name]:
continue
# Both set to different values: the kwarg wins (setdefault below is a no-op). Say so explicitly rather
# than implying the deprecated field was migrated, which would mislead — the field value is discarded.
if kwarg_name in kwargs:
warnings.warn(
f"{field_name}={data[field_name]!r} is ignored because lr_scheduler_kwargs already sets "
f"{kwarg_name!r}={kwargs[kwarg_name]!r} (the kwarg wins); remove the deprecated {field_name}.",
FutureWarning,
stacklevel=2,
)
continue
warnings.warn(
f"{field_name} is deprecated; pass it via lr_scheduler_kwargs={{'{kwarg_name}': ...}} instead.",
FutureWarning,
stacklevel=2,
)
kwargs[kwarg_name] = data[field_name]
if kwargs:
data["lr_scheduler_kwargs"] = kwargs
return data
@model_validator(mode="before")
@classmethod
def _desugar_callable_lr_scheduler(cls, data: Any) -> Any:
"""Desugar a reconstructable callable lr_scheduler into its serializable string form.
A callable ``lr_scheduler`` (a class or ``functools.partial``) that can be imported is rewritten to a dotted
import path plus ``lr_scheduler_kwargs`` so the config round-trips through ``training_config.json``. User-
supplied ``lr_scheduler_kwargs`` are ignored for callable schedulers (bake arguments into the callable instead).
Non-reconstructable callables are kept as-is and only warned about.
"""
if not isinstance(data, dict):
return data
lr_scheduler = data.get("lr_scheduler")
if lr_scheduler is None or isinstance(lr_scheduler, str) or not callable(lr_scheduler):
return data
if data.get("lr_scheduler_kwargs"):
warnings.warn(
"lr_scheduler_kwargs is ignored when lr_scheduler is a callable; bake arguments into the "
"callable (for example with functools.partial) instead.",
stacklevel=2,
)
path, kwargs, reason = _desugar_scheduler_callable(lr_scheduler)
if reason is None:
data["lr_scheduler"] = path
data["lr_scheduler_kwargs"] = kwargs
else:
data["lr_scheduler_kwargs"] = {}
label = getattr(lr_scheduler, "__qualname__", None) or repr(lr_scheduler)
warnings.warn(
f"lr_scheduler callable {label!r} cannot be saved to training_config.json and restored: "
f"{reason}. Training proceeds with the in-memory callable; only saved-config "
"reproducibility is affected.",
stacklevel=2,
)
return data
@field_validator("lr_scheduler", mode="after")
@classmethod
def validate_lr_scheduler_name(cls, v: str | Callable[..., SchedulerType]) -> str | Callable[..., SchedulerType]:
"""Validate a string lr_scheduler: a bare name must be a managed preset, else use a dotted path."""
if not isinstance(v, str):
return v
lr_scheduler = v.strip()
if not lr_scheduler:
raise ValueError("lr_scheduler must be a non-empty string.")
# Bare names must be a managed preset; dotted import paths are validated lazily at train start.
if "." not in lr_scheduler and not _is_managed_scheduler_name(lr_scheduler):
presets = ", ".join(sorted(_MANAGED_SCHEDULER_PRESETS))
raise ValueError(
f"Unknown lr_scheduler {v!r}. Bare names must be a managed preset ({presets}); "
"use a full dotted import path (e.g. 'torch.optim.lr_scheduler.StepLR') or a callable "
"for anything else."
)
return lr_scheduler
@field_validator("prefetch_factor", mode="after")
@classmethod
def validate_prefetch_factor(cls, v: int | None) -> int | None:
"""Validate prefetch_factor is None or >= 1."""
if v is not None and v < 1:
raise ValueError("prefetch_factor must be >= 1 when provided.")
return v
@field_validator("dataset_dir", "output_dir", mode="before")
@classmethod
def expand_paths(cls, v: PathLikeStr | None) -> str | None:
"""Expand and normalize dataset/output directory paths via ``os.fspath`` → ``expanduser`` → ``realpath``."""
if v is None:
return v
return os.path.realpath(os.path.expanduser(os.fspath(v)))
@field_validator("resume", mode="before")
@classmethod
def _coerce_resume_path(cls, v: PathLikeStr | None) -> str | None:
"""Normalise the resume checkpoint value to ``str`` without resolving it.
Unlike ``dataset_dir``/``output_dir``, ``resume`` is forwarded verbatim to PyTorch Lightning's
``trainer.fit(ckpt_path=...)``, which also accepts sentinel values such as ``"last"``. Running
``os.path.realpath`` would rewrite those sentinels into spurious absolute paths, so this validator only coerces
the type (``Path`` -> ``str``) and leaves the value untouched.
"""
if v is None:
return v
return os.fspath(v)