Training API

Training runs on PyTorch Lightning. run_training is the entry point the yolonas train CLI calls; the pieces below are what it wires together, and each can be used on its own if you would rather drive Lightning yourself.

Entry point

modern_yolonas.training.run.run_training(model_name, recipe, train_dataset, val_dataset, val_ann_file=None, output_dir='runs/train', num_classes=80, pretrained=True, devices='auto', logger='csv', resume_path=None, epochs=None, batch_size=None)

Run a full training session, returning path to the best checkpoint.

Parameters:

Name	Type	Description	Default
model_name	str	One of 'yolo_nas_s', 'yolo_nas_m', 'yolo_nas_l'.	required
recipe	dict	Recipe dict (from recipes.py).	required
train_dataset	Dataset	Training dataset.	required
val_dataset	Dataset | None	Validation dataset.	required
val_ann_file	str | Path | None	COCO annotation file for mAP evaluation.	None
output_dir	str | Path	Directory for checkpoints and logs.	'runs/train'
num_classes	int	Number of object classes.	80
pretrained	bool	Whether to start from pretrained COCO weights.	True
devices	str | int | list[int]	Lightning devices spec.	'auto'
logger	str	Logger backend ('csv', 'tensorboard', 'wandb').	'csv'
resume_path	str | Path | None	Optional checkpoint to resume from.	None
epochs	int | None	Override recipe epochs.	None
batch_size	int | None	Override recipe batch_size.	None

Returns:

Type	Description
Path	Path to the best checkpoint file.

Source code in src/modern_yolonas/training/run.py

def run_training(
    model_name: str,
    recipe: dict,
    train_dataset: Dataset,
    val_dataset: Dataset | None,
    val_ann_file: str | Path | None = None,
    output_dir: str | Path = "runs/train",
    num_classes: int = 80,
    pretrained: bool = True,
    devices: str | int | list[int] = "auto",
    logger: str = "csv",
    resume_path: str | Path | None = None,
    epochs: int | None = None,
    batch_size: int | None = None,
) -> Path:
    """Run a full training session, returning path to the best checkpoint.

    Args:
        model_name: One of 'yolo_nas_s', 'yolo_nas_m', 'yolo_nas_l'.
        recipe: Recipe dict (from recipes.py).
        train_dataset: Training dataset.
        val_dataset: Validation dataset.
        val_ann_file: COCO annotation file for mAP evaluation.
        output_dir: Directory for checkpoints and logs.
        num_classes: Number of object classes.
        pretrained: Whether to start from pretrained COCO weights.
        devices: Lightning devices spec.
        logger: Logger backend ('csv', 'tensorboard', 'wandb').
        resume_path: Optional checkpoint to resume from.
        epochs: Override recipe epochs.
        batch_size: Override recipe batch_size.

    Returns:
        Path to the best checkpoint file.
    """
    output_dir = Path(output_dir)
    epochs = epochs or recipe["epochs"]
    batch_size = batch_size or recipe["batch_size"]

    # Build model
    builder = MODEL_BUILDERS[model_name]
    model = builder(pretrained=pretrained, num_classes=num_classes)

    # Warmup steps
    warmup_epochs = recipe.get("warmup_epochs", 3)
    steps_per_epoch = max(1, len(train_dataset) // batch_size)
    warmup_steps = warmup_epochs * steps_per_epoch

    # Lightning module
    lit_model = YoloNASLightningModule(
        channels_last=recipe.get("channels_last", True),
        model=model,
        num_classes=num_classes,
        lr=recipe["lr"],
        optimizer_name=recipe["optimizer"],
        weight_decay=recipe["weight_decay"],
        warmup_steps=warmup_steps,
        cosine_final_lr_ratio=recipe["cosine_final_lr_ratio"],
        val_ann_file=val_ann_file,
        conf_threshold=recipe.get("conf_threshold", 0.001),
        iou_threshold=recipe.get("iou_threshold", 0.65),
        input_size=recipe["input_size"],
    )

    # Data module
    data_module = DetectionDataModule(
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        batch_size=batch_size,
        num_workers=recipe.get("workers", 8),
    )

    # Logger
    if logger == "csv":
        logger_instance = L.pytorch.loggers.CSVLogger(str(output_dir))
    elif logger == "tensorboard":
        logger_instance = L.pytorch.loggers.TensorBoardLogger(str(output_dir))
    else:
        logger_instance = L.pytorch.loggers.WandbLogger(project="yolonas", save_dir=str(output_dir))

    # Callbacks
    monitor = "val/mAP" if val_ann_file else "val/loss"
    mode = "max" if val_ann_file else "min"
    checkpoint_cb = ModelCheckpoint(
        dirpath=str(output_dir / "checkpoints"),
        monitor=monitor,
        mode=mode,
        save_last=True,
        filename="best-{epoch:02d}-{" + monitor.replace("/", "_") + ":.4f}",
    )
    callbacks = [
        EMACallback(decay=recipe.get("ema_decay", 0.9997)),
        checkpoint_cb,
    ]

    # CloseMosaicCallback
    aug = recipe.get("augmentations", {})
    close_mosaic_epochs = aug.get("close_mosaic_epochs", 0)
    if close_mosaic_epochs > 0 and (aug.get("mosaic") or aug.get("mixup")):
        callbacks.append(CloseMosaicCallback(close_mosaic_epochs=close_mosaic_epochs))

    # Trainer
    trainer = L.Trainer(
        max_epochs=epochs,
        accelerator="auto",
        devices=devices,
        strategy="auto",
        precision=recipe.get("precision", "16-mixed"),
        # From-scratch training with AMP NaNs in the first epoch without this.
        gradient_clip_val=recipe.get("grad_clip", 10.0) or None,
        callbacks=callbacks,
        logger=logger_instance,
        default_root_dir=str(output_dir),
    )

    if resume_path:
        trainer.fit(lit_model, datamodule=data_module, ckpt_path=str(resume_path))
    else:
        trainer.fit(lit_model, datamodule=data_module)

    best_path = checkpoint_cb.best_model_path
    return Path(best_path) if best_path else output_dir / "checkpoints" / "last.ckpt"

Lightning module and data

modern_yolonas.training.lightning_module.YoloNASLightningModule

Bases: LightningModule

Lightning module wrapping a YOLO-NAS model for training.

Parameters:

Name	Type	Description	Default
model	Module	YoloNAS or QuantizedYoloNAS model.	required
num_classes	int	Number of object classes.	80
lr	float	Learning rate.	0.0002
optimizer_name	str	Optimizer name ('adamw' or 'sgd').	'adamw'
weight_decay	float	Weight decay.	1e-05
warmup_steps	int	LR warmup steps.	1000
cosine_final_lr_ratio	float	Final LR as fraction of initial LR.	0.1
val_ann_file	str | Path | None	Path to COCO annotation JSON for mAP evaluation. When set, validation computes mAP instead of loss.	None
val_dataset_ids	list[int] | None	List of COCO image IDs matching validation dataset order. If None, extracted from the val dataloader's dataset.	None
conf_threshold	float	Confidence threshold for mAP evaluation.	0.001
iou_threshold	float	NMS IoU threshold for mAP evaluation.	0.65
input_size	int	Square size the images are letterboxed to. The evaluator needs it to map predictions back to original image coordinates before comparing them with the ground truth.	640
channels_last	bool	Hold activations in NHWC. Ampere and newer run convolutions under AMP faster in that layout — measured +15% on a 3060 Laptop, 2026-09-19 — and it changes memory layout only, not results.	False

Source code in src/modern_yolonas/training/lightning_module.py

class YoloNASLightningModule(L.LightningModule):
    """Lightning module wrapping a YOLO-NAS model for training.

    Args:
        model: YoloNAS or QuantizedYoloNAS model.
        num_classes: Number of object classes.
        lr: Learning rate.
        optimizer_name: Optimizer name ('adamw' or 'sgd').
        weight_decay: Weight decay.
        warmup_steps: LR warmup steps.
        cosine_final_lr_ratio: Final LR as fraction of initial LR.
        val_ann_file: Path to COCO annotation JSON for mAP evaluation.
            When set, validation computes mAP instead of loss.
        val_dataset_ids: List of COCO image IDs matching validation dataset order.
            If None, extracted from the val dataloader's dataset.
        conf_threshold: Confidence threshold for mAP evaluation.
        iou_threshold: NMS IoU threshold for mAP evaluation.
        input_size: Square size the images are letterboxed to. The evaluator needs it to
            map predictions back to original image coordinates before comparing them
            with the ground truth.
        channels_last: Hold activations in NHWC. Ampere and newer run convolutions
            under AMP faster in that layout — measured +15% on a 3060 Laptop,
            2026-09-19 — and it changes memory layout only, not results.
    """

    def __init__(
        self,
        model: nn.Module,
        num_classes: int = 80,
        lr: float = 2e-4,
        optimizer_name: str = "adamw",
        weight_decay: float = 1e-5,
        warmup_steps: int = 1000,
        cosine_final_lr_ratio: float = 0.1,
        val_ann_file: str | Path | None = None,
        val_dataset_ids: list[int] | None = None,
        conf_threshold: float = 0.001,
        iou_threshold: float = 0.65,
        input_size: int = 640,
        channels_last: bool = False,
    ):
        super().__init__()
        self.model = model.to(memory_format=torch.channels_last) if channels_last else model
        self.channels_last = channels_last
        self.criterion = PPYoloELoss(num_classes=num_classes)
        self.val_ann_file = val_ann_file
        self.val_dataset_ids = val_dataset_ids
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.input_size = input_size
        self._evaluator = None
        self.save_hyperparameters(ignore=["model"])

    def forward(self, x):
        return self.model(x)

    def on_after_batch_transfer(self, batch, dataloader_idx):
        """Finish preprocessing on the device, after the batch has crossed the bus.

        ``Normalize(dtype="uint8")`` leaves the divide-by-255 to this hook so that a
        quarter as many bytes cross ``pin_memory`` and PCIe. The check is on the dtype
        rather than on a flag, so a float32 pipeline passes through untouched and
        either kind of dataloader works against either kind of module.
        """
        images, targets = batch
        if images.dtype == torch.uint8:
            images = images.float().div_(255.0)
        if self.channels_last:
            images = images.contiguous(memory_format=torch.channels_last)
        return images, targets

    def training_step(self, batch, batch_idx):
        images, targets = batch
        predictions = self.model(images)
        loss, loss_dict = self.criterion(predictions, targets)

        self.log("train/loss", loss, prog_bar=True)
        self.log("train/cls_loss", loss_dict["cls_loss"])
        self.log("train/iou_loss", loss_dict["iou_loss"])
        self.log("train/dfl_loss", loss_dict["dfl_loss"])

        if self._trainer is not None:
            schedulers = self.lr_schedulers()
            if schedulers is not None:
                lr = schedulers.get_last_lr()[0]
                self.log("train/lr", lr, prog_bar=True)

        return loss

    def on_validation_epoch_start(self):
        if self.val_ann_file is not None:
            from modern_yolonas.training.metrics import COCOEvaluator

            self._evaluator = COCOEvaluator(self.val_ann_file, input_size=self.input_size)
            # Cache image IDs from the val dataset if not provided
            if self.val_dataset_ids is None and self._trainer is not None:
                val_dl = self.trainer.val_dataloaders
                if val_dl is not None and hasattr(val_dl.dataset, "ids"):
                    self.val_dataset_ids = val_dl.dataset.ids

    def validation_step(self, batch, batch_idx):
        images, targets = batch

        # Model is in eval mode during validation (set by Lightning).
        # In eval mode, NDFLHeads returns (bboxes, scores) directly.
        # In train mode, it returns ((decoded, raw)).
        if self._evaluator is not None:
            from modern_yolonas.inference.postprocess import postprocess

            pred_bboxes, pred_scores = self.model(images)
            results = postprocess(
                pred_bboxes, pred_scores,
                conf_threshold=self.conf_threshold,
                iou_threshold=self.iou_threshold,
            )

            # Resolve image IDs for this batch
            batch_size = images.shape[0]
            start_idx = batch_idx * batch_size
            if self.val_dataset_ids is not None:
                image_ids = [
                    self.val_dataset_ids[start_idx + i]
                    for i in range(batch_size)
                    if start_idx + i < len(self.val_dataset_ids)
                ]
            else:
                image_ids = list(range(start_idx, start_idx + batch_size))

            boxes_list = [r[0] for r in results]
            scores_list = [r[1] for r in results]
            class_ids_list = [r[2] for r in results]
            self._evaluator.update(image_ids, boxes_list, scores_list, class_ids_list)
        else:
            # The loss needs the raw, undecoded predictions, which the head normally
            # returns only in training mode. Flipping the model to train() to get them
            # would let BatchNorm overwrite its running statistics with validation-set
            # statistics — `torch.no_grad` stops gradients, not buffer updates — and
            # then persist them into the checkpoint. `return_raw_outputs` changes the
            # return value without touching any module's training flag.
            heads = getattr(self.model, "heads", None)
            if heads is None:
                raise AttributeError(
                    "validation loss needs raw predictions from NDFLHeads, but the model "
                    "has no `heads` attribute"
                )
            heads.return_raw_outputs = True
            try:
                predictions = self.model(images)
                loss, loss_dict = self.criterion(predictions, targets)
            finally:
                heads.return_raw_outputs = False

            self.log("val/loss", loss, prog_bar=True, sync_dist=True)
            self.log("val/cls_loss", loss_dict["cls_loss"], sync_dist=True)
            self.log("val/iou_loss", loss_dict["iou_loss"], sync_dist=True)
            self.log("val/dfl_loss", loss_dict["dfl_loss"], sync_dist=True)

    def on_validation_epoch_end(self):
        if self._evaluator is not None:
            metrics = self._evaluator.evaluate()
            self.log("val/mAP", metrics["mAP"], prog_bar=True, sync_dist=True)
            self.log("val/mAP_50", metrics["mAP_50"], sync_dist=True)
            self.log("val/mAP_75", metrics["mAP_75"], sync_dist=True)
            self._evaluator = None

    def configure_optimizers(self):
        optimizer = create_optimizer(
            self.model,
            name=self.hparams.optimizer_name,
            lr=self.hparams.lr,
            weight_decay=self.hparams.weight_decay,
        )
        total_steps = self.trainer.estimated_stepping_batches
        scheduler = cosine_with_warmup(
            optimizer,
            warmup_steps=self.hparams.warmup_steps,
            total_steps=total_steps,
            cosine_final_lr_ratio=self.hparams.cosine_final_lr_ratio,
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {"scheduler": scheduler, "interval": "step"},
        }

on_after_batch_transfer(batch, dataloader_idx)

Finish preprocessing on the device, after the batch has crossed the bus.

Normalize(dtype="uint8") leaves the divide-by-255 to this hook so that a quarter as many bytes cross pin_memory and PCIe. The check is on the dtype rather than on a flag, so a float32 pipeline passes through untouched and either kind of dataloader works against either kind of module.

Source code in src/modern_yolonas/training/lightning_module.py

def on_after_batch_transfer(self, batch, dataloader_idx):
    """Finish preprocessing on the device, after the batch has crossed the bus.

    ``Normalize(dtype="uint8")`` leaves the divide-by-255 to this hook so that a
    quarter as many bytes cross ``pin_memory`` and PCIe. The check is on the dtype
    rather than on a flag, so a float32 pipeline passes through untouched and
    either kind of dataloader works against either kind of module.
    """
    images, targets = batch
    if images.dtype == torch.uint8:
        images = images.float().div_(255.0)
    if self.channels_last:
        images = images.contiguous(memory_format=torch.channels_last)
    return images, targets

modern_yolonas.training.data_module.DetectionDataModule

Bases: LightningDataModule

DataModule wrapping training and validation detection datasets.

Parameters:

Name	Type	Description	Default
train_dataset	Dataset	Training dataset.	required
val_dataset	Dataset | None	Optional validation dataset.	None
batch_size	int	Batch size per device.	32
num_workers	int	DataLoader workers.	8

Source code in src/modern_yolonas/training/data_module.py

class DetectionDataModule(L.LightningDataModule):
    """DataModule wrapping training and validation detection datasets.

    Args:
        train_dataset: Training dataset.
        val_dataset: Optional validation dataset.
        batch_size: Batch size per device.
        num_workers: DataLoader workers.
    """

    def __init__(
        self,
        train_dataset: Dataset,
        val_dataset: Dataset | None = None,
        batch_size: int = 32,
        num_workers: int = 8,
    ):
        super().__init__()
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.batch_size = batch_size
        self.num_workers = num_workers

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            collate_fn=detection_collate_fn,
            pin_memory=True,
            drop_last=True,
        )

    def val_dataloader(self):
        if self.val_dataset is None:
            return None
        return DataLoader(
            self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            collate_fn=detection_collate_fn,
            pin_memory=True,
        )

Loss and metrics

modern_yolonas.training.loss.PPYoloELoss

Bases: Module

Combined loss for YOLO-NAS training.

Components: - VarifocalLoss (classification) - GIoULoss (box regression, weighted by assigned scores) - DFLLoss (distribution focal loss, weighted by assigned scores)

Weighted sum: cls_weight * vfl + iou_weight * giou + dfl_weight * dfl All terms normalized by assigned_scores_sum (matching super-gradients).

Parameters:

Name	Type	Description	Default
num_classes	int	Number of object classes.	80
reg_max	int	Distribution regression maximum.	16
cls_weight	float	Classification loss weight.	1.0
iou_weight	float	Box regression loss weight.	2.5
dfl_weight	float	Distribution focal loss weight.	0.5
static_assigner_epochs	int	Use ATSS for the first N epochs (0 to disable).	4

Source code in src/modern_yolonas/training/loss.py

class PPYoloELoss(nn.Module):
    """Combined loss for YOLO-NAS training.

    Components:
    - VarifocalLoss (classification)
    - GIoULoss (box regression, weighted by assigned scores)
    - DFLLoss (distribution focal loss, weighted by assigned scores)

    Weighted sum: cls_weight * vfl + iou_weight * giou + dfl_weight * dfl
    All terms normalized by ``assigned_scores_sum`` (matching super-gradients).

    Args:
        num_classes: Number of object classes.
        reg_max: Distribution regression maximum.
        cls_weight: Classification loss weight.
        iou_weight: Box regression loss weight.
        dfl_weight: Distribution focal loss weight.
        static_assigner_epochs: Use ATSS for the first N epochs (0 to disable).
    """

    def __init__(
        self,
        num_classes: int = 80,
        reg_max: int = 16,
        cls_weight: float = 1.0,
        iou_weight: float = 2.5,
        dfl_weight: float = 0.5,
        static_assigner_epochs: int = 4,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.reg_max = reg_max
        self.cls_weight = cls_weight
        self.iou_weight = iou_weight
        self.dfl_weight = dfl_weight
        self.static_assigner_epochs = static_assigner_epochs

        self.assigner = TaskAlignedAssigner()
        self.static_assigner = ATSSAssigner(topk=9) if static_assigner_epochs > 0 else None
        self.vfl = VarifocalLoss()
        self.giou_loss = GIoULoss()
        self.dfl_loss = DFLLoss()

    @staticmethod
    def _crowd_ignore_mask(
        crowd_targets: Tensor,
        anchor_points: Tensor,
        batch_size: int,
        img_w: Tensor | float,
        img_h: Tensor | float,
    ) -> Tensor | None:
        """``[B, N]`` anchors whose centre lies inside any crowd box.

        ``anchor_points`` are in image pixels, the same space the ground-truth boxes
        are scaled into, so the containment test needs no stride correction.

        Returns ``None`` when there are no crowd regions, so the common path adds
        nothing.
        """
        if crowd_targets.numel() == 0:
            return None

        device = anchor_points.device
        mask = torch.zeros(batch_size, anchor_points.shape[0], dtype=torch.bool, device=device)
        for b in range(batch_size):
            rows = crowd_targets[crowd_targets[:, 0] == b]
            if rows.numel() == 0:
                continue
            xc, yc = rows[:, 2] * img_w, rows[:, 3] * img_h
            w, h = rows[:, 4] * img_w, rows[:, 5] * img_h
            x1, y1 = (xc - w / 2).view(1, -1), (yc - h / 2).view(1, -1)
            x2, y2 = (xc + w / 2).view(1, -1), (yc + h / 2).view(1, -1)
            px, py = anchor_points[:, 0:1], anchor_points[:, 1:2]
            mask[b] = ((px >= x1) & (px <= x2) & (py >= y1) & (py <= y2)).any(dim=1)
        return mask

    def _bbox2dist(self, anchor_points: Tensor, gt_bboxes: Tensor) -> Tensor:
        """Convert bounding boxes to distances from anchor points."""
        x1y1 = anchor_points - gt_bboxes[..., :2]
        x2y2 = gt_bboxes[..., 2:] - anchor_points
        dist = torch.cat([x1y1, x2y2], dim=-1)
        return dist.clamp(0, self.reg_max - 0.01)

    def forward(
        self,
        predictions: tuple,
        targets: Tensor,
        input_size: tuple[int, int] | None = None,
        epoch: int | None = None,
    ) -> tuple[Tensor, dict[str, float]]:
        """Compute loss.

        Args:
            predictions: ``(decoded_predictions, raw_predictions)`` from NDFLHeads in training mode.
                decoded_predictions: ``(pred_bboxes [B,N,4], pred_scores [B,N,C])``
                raw_predictions: ``(cls_logits [B,N,C], reg_distri [B,N,4*(reg_max+1)],
                                   anchors, anchor_points, num_anchors_list, stride_tensor)``
            targets: ``[sum(N_i), 6]`` with ``[batch_idx, class_id, x, y, w, h]`` (normalized xywh).
            input_size: ``(H, W)`` of the input image. If None, inferred from anchor grid.
            epoch: Current training epoch (used for ATSS warmup).

        Returns:
            (total_loss, loss_dict)
        """
        (pred_bboxes_decoded, pred_scores_decoded), (
            cls_logits,
            reg_distri,
            anchors,
            anchor_points,
            num_anchors_list,
            stride_tensor,
        ) = predictions

        batch_size = cls_logits.shape[0]
        device = cls_logits.device

        # Determine input image size for scaling normalized GT to pixel coords
        if input_size is not None:
            img_h, img_w = input_size[0], input_size[1]
        else:
            inferred = (anchor_points.max(dim=0).values + stride_tensor.min() / 2).clamp(min=1)
            img_w, img_h = inferred[0], inferred[1]

        # Validate GT class labels
        if targets.numel() > 0:
            class_ids = targets[:, 1]
            if (class_ids < CROWD_CLASS).any() or (class_ids >= self.num_classes).any():
                logger.warning(
                    "GT class labels out of range [0, %d): min=%d, max=%d. "
                    "Check your dataset labels.",
                    self.num_classes,
                    int(class_ids.min().item()),
                    int(class_ids.max().item()),
                )

        # Crowd regions are not trainable targets. They are split out here and used
        # only to build an ignore mask, so the anchors they cover become neither
        # positives nor background.
        is_crowd = targets[:, 1] == CROWD_CLASS
        crowd_targets = targets[is_crowd]
        targets = targets[~is_crowd]

        # Prepare GT in format expected by assigner
        gt_labels_list = []
        gt_bboxes_list = []
        for b in range(batch_size):
            mask = targets[:, 0] == b
            if mask.any():
                t = targets[mask]
                gt_labels_list.append(t[:, 1:2])
                xc, yc, w, h = t[:, 2], t[:, 3], t[:, 4], t[:, 5]
                xc, w = xc * img_w, w * img_w
                yc, h = yc * img_h, h * img_h
                gt_bboxes_list.append(torch.stack([
                    xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2
                ], dim=-1))
            else:
                gt_labels_list.append(torch.zeros(0, 1, device=device))
                gt_bboxes_list.append(torch.zeros(0, 4, device=device))

        max_gt = max(len(g) for g in gt_labels_list)
        if max_gt == 0:
            zero_loss = cls_logits.sum() * 0.0
            return zero_loss, {"cls_loss": 0.0, "iou_loss": 0.0, "dfl_loss": 0.0, "total_loss": 0.0}

        gt_labels = torch.zeros(batch_size, max_gt, 1, device=device)
        gt_bboxes = torch.zeros(batch_size, max_gt, 4, device=device)
        mask_gt = torch.zeros(batch_size, max_gt, 1, device=device)

        for b in range(batch_size):
            n = len(gt_labels_list[b])
            if n > 0:
                gt_labels[b, :n] = gt_labels_list[b]
                gt_bboxes[b, :n] = gt_bboxes_list[b]
                mask_gt[b, :n] = 1.0

        # Select assigner: ATSS for warmup, TAL after
        use_static = (
            self.static_assigner is not None
            and epoch is not None
            and epoch < self.static_assigner_epochs
        )

        if use_static:
            assigned_labels, assigned_bboxes, assigned_scores, fg_mask = self.static_assigner.assign(
                anchors, num_anchors_list, gt_labels, gt_bboxes, mask_gt, self.num_classes,
                pred_bboxes=pred_bboxes_decoded.detach(),
            )
        else:
            assigned_labels, assigned_bboxes, assigned_scores, fg_mask = self.assigner.assign(
                pred_scores_decoded, pred_bboxes_decoded, anchor_points,
                gt_labels, gt_bboxes, mask_gt,
            )

        # Normalization: sum of assigned soft scores (matches super-gradients)
        assigned_scores_sum = assigned_scores.sum().clamp(min=1)

        # Anchors whose centre falls in a crowd region, and that the assigner did not
        # already claim for a real object. A real instance overlapping a crowd box
        # keeps its positive; only the leftover background anchors are dropped.
        ignore_mask = self._crowd_ignore_mask(
            crowd_targets, anchor_points, batch_size, img_w, img_h
        )
        if ignore_mask is not None:
            ignore_mask = ignore_mask & ~fg_mask.bool()

        # Classification loss (VFL)
        cls_loss = self.vfl(
            cls_logits, assigned_scores, (assigned_scores > 0).float(), ignore_mask
        ) / assigned_scores_sum

        # Box regression loss (GIoU) — weighted by per-anchor assigned scores
        if fg_mask.any():
            pos_pred_bboxes = pred_bboxes_decoded[fg_mask]
            pos_target_bboxes = assigned_bboxes[fg_mask]
            bbox_weight = assigned_scores[fg_mask].sum(-1)  # [num_pos]
            iou_loss_per_anchor = self.giou_loss(pos_pred_bboxes, pos_target_bboxes)  # [num_pos]
            iou_loss = (iou_loss_per_anchor * bbox_weight).sum() / assigned_scores_sum
        else:
            iou_loss = torch.tensor(0.0, device=device)

        # DFL loss — weighted by per-anchor assigned scores
        if fg_mask.any():
            pos_reg_distri = reg_distri[fg_mask]
            pos_anchor_points = anchor_points.unsqueeze(0).expand(batch_size, -1, -1)[fg_mask]
            pos_stride = stride_tensor.unsqueeze(0).expand(batch_size, -1, -1)[fg_mask]
            pos_target_bboxes = assigned_bboxes[fg_mask] / pos_stride
            pos_anchor_points_scaled = pos_anchor_points / pos_stride
            target_dist = self._bbox2dist(pos_anchor_points_scaled, pos_target_bboxes)
            dfl_loss_per_anchor = self.dfl_loss(pos_reg_distri, target_dist)  # [num_pos]
            dfl_loss = (dfl_loss_per_anchor * bbox_weight).sum() / assigned_scores_sum
        else:
            dfl_loss = torch.tensor(0.0, device=device)

        total_loss = self.cls_weight * cls_loss + self.iou_weight * iou_loss + self.dfl_weight * dfl_loss

        loss_dict = {
            "cls_loss": cls_loss.item(),
            "iou_loss": iou_loss.item(),
            "dfl_loss": dfl_loss.item(),
            "total_loss": total_loss.item(),
        }

        return total_loss, loss_dict

forward(predictions, targets, input_size=None, epoch=None)

Compute loss.

Parameters:

Name	Type	Description	Default
predictions	tuple	(decoded_predictions, raw_predictions) from NDFLHeads in training mode. decoded_predictions: (pred_bboxes [B,N,4], pred_scores [B,N,C]) raw_predictions: (cls_logits [B,N,C], reg_distri [B,N,4*(reg_max+1)], anchors, anchor_points, num_anchors_list, stride_tensor)	required
targets	Tensor	[sum(N_i), 6] with [batch_idx, class_id, x, y, w, h] (normalized xywh).	required
input_size	tuple[int, int] | None	(H, W) of the input image. If None, inferred from anchor grid.	None
epoch	int | None	Current training epoch (used for ATSS warmup).	None

Returns:

Type	Description
tuple[Tensor, dict[str, float]]	(total_loss, loss_dict)

Source code in src/modern_yolonas/training/loss.py

def forward(
    self,
    predictions: tuple,
    targets: Tensor,
    input_size: tuple[int, int] | None = None,
    epoch: int | None = None,
) -> tuple[Tensor, dict[str, float]]:
    """Compute loss.

    Args:
        predictions: ``(decoded_predictions, raw_predictions)`` from NDFLHeads in training mode.
            decoded_predictions: ``(pred_bboxes [B,N,4], pred_scores [B,N,C])``
            raw_predictions: ``(cls_logits [B,N,C], reg_distri [B,N,4*(reg_max+1)],
                               anchors, anchor_points, num_anchors_list, stride_tensor)``
        targets: ``[sum(N_i), 6]`` with ``[batch_idx, class_id, x, y, w, h]`` (normalized xywh).
        input_size: ``(H, W)`` of the input image. If None, inferred from anchor grid.
        epoch: Current training epoch (used for ATSS warmup).

    Returns:
        (total_loss, loss_dict)
    """
    (pred_bboxes_decoded, pred_scores_decoded), (
        cls_logits,
        reg_distri,
        anchors,
        anchor_points,
        num_anchors_list,
        stride_tensor,
    ) = predictions

    batch_size = cls_logits.shape[0]
    device = cls_logits.device

    # Determine input image size for scaling normalized GT to pixel coords
    if input_size is not None:
        img_h, img_w = input_size[0], input_size[1]
    else:
        inferred = (anchor_points.max(dim=0).values + stride_tensor.min() / 2).clamp(min=1)
        img_w, img_h = inferred[0], inferred[1]

    # Validate GT class labels
    if targets.numel() > 0:
        class_ids = targets[:, 1]
        if (class_ids < CROWD_CLASS).any() or (class_ids >= self.num_classes).any():
            logger.warning(
                "GT class labels out of range [0, %d): min=%d, max=%d. "
                "Check your dataset labels.",
                self.num_classes,
                int(class_ids.min().item()),
                int(class_ids.max().item()),
            )

    # Crowd regions are not trainable targets. They are split out here and used
    # only to build an ignore mask, so the anchors they cover become neither
    # positives nor background.
    is_crowd = targets[:, 1] == CROWD_CLASS
    crowd_targets = targets[is_crowd]
    targets = targets[~is_crowd]

    # Prepare GT in format expected by assigner
    gt_labels_list = []
    gt_bboxes_list = []
    for b in range(batch_size):
        mask = targets[:, 0] == b
        if mask.any():
            t = targets[mask]
            gt_labels_list.append(t[:, 1:2])
            xc, yc, w, h = t[:, 2], t[:, 3], t[:, 4], t[:, 5]
            xc, w = xc * img_w, w * img_w
            yc, h = yc * img_h, h * img_h
            gt_bboxes_list.append(torch.stack([
                xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2
            ], dim=-1))
        else:
            gt_labels_list.append(torch.zeros(0, 1, device=device))
            gt_bboxes_list.append(torch.zeros(0, 4, device=device))

    max_gt = max(len(g) for g in gt_labels_list)
    if max_gt == 0:
        zero_loss = cls_logits.sum() * 0.0
        return zero_loss, {"cls_loss": 0.0, "iou_loss": 0.0, "dfl_loss": 0.0, "total_loss": 0.0}

    gt_labels = torch.zeros(batch_size, max_gt, 1, device=device)
    gt_bboxes = torch.zeros(batch_size, max_gt, 4, device=device)
    mask_gt = torch.zeros(batch_size, max_gt, 1, device=device)

    for b in range(batch_size):
        n = len(gt_labels_list[b])
        if n > 0:
            gt_labels[b, :n] = gt_labels_list[b]
            gt_bboxes[b, :n] = gt_bboxes_list[b]
            mask_gt[b, :n] = 1.0

    # Select assigner: ATSS for warmup, TAL after
    use_static = (
        self.static_assigner is not None
        and epoch is not None
        and epoch < self.static_assigner_epochs
    )

    if use_static:
        assigned_labels, assigned_bboxes, assigned_scores, fg_mask = self.static_assigner.assign(
            anchors, num_anchors_list, gt_labels, gt_bboxes, mask_gt, self.num_classes,
            pred_bboxes=pred_bboxes_decoded.detach(),
        )
    else:
        assigned_labels, assigned_bboxes, assigned_scores, fg_mask = self.assigner.assign(
            pred_scores_decoded, pred_bboxes_decoded, anchor_points,
            gt_labels, gt_bboxes, mask_gt,
        )

    # Normalization: sum of assigned soft scores (matches super-gradients)
    assigned_scores_sum = assigned_scores.sum().clamp(min=1)

    # Anchors whose centre falls in a crowd region, and that the assigner did not
    # already claim for a real object. A real instance overlapping a crowd box
    # keeps its positive; only the leftover background anchors are dropped.
    ignore_mask = self._crowd_ignore_mask(
        crowd_targets, anchor_points, batch_size, img_w, img_h
    )
    if ignore_mask is not None:
        ignore_mask = ignore_mask & ~fg_mask.bool()

    # Classification loss (VFL)
    cls_loss = self.vfl(
        cls_logits, assigned_scores, (assigned_scores > 0).float(), ignore_mask
    ) / assigned_scores_sum

    # Box regression loss (GIoU) — weighted by per-anchor assigned scores
    if fg_mask.any():
        pos_pred_bboxes = pred_bboxes_decoded[fg_mask]
        pos_target_bboxes = assigned_bboxes[fg_mask]
        bbox_weight = assigned_scores[fg_mask].sum(-1)  # [num_pos]
        iou_loss_per_anchor = self.giou_loss(pos_pred_bboxes, pos_target_bboxes)  # [num_pos]
        iou_loss = (iou_loss_per_anchor * bbox_weight).sum() / assigned_scores_sum
    else:
        iou_loss = torch.tensor(0.0, device=device)

    # DFL loss — weighted by per-anchor assigned scores
    if fg_mask.any():
        pos_reg_distri = reg_distri[fg_mask]
        pos_anchor_points = anchor_points.unsqueeze(0).expand(batch_size, -1, -1)[fg_mask]
        pos_stride = stride_tensor.unsqueeze(0).expand(batch_size, -1, -1)[fg_mask]
        pos_target_bboxes = assigned_bboxes[fg_mask] / pos_stride
        pos_anchor_points_scaled = pos_anchor_points / pos_stride
        target_dist = self._bbox2dist(pos_anchor_points_scaled, pos_target_bboxes)
        dfl_loss_per_anchor = self.dfl_loss(pos_reg_distri, target_dist)  # [num_pos]
        dfl_loss = (dfl_loss_per_anchor * bbox_weight).sum() / assigned_scores_sum
    else:
        dfl_loss = torch.tensor(0.0, device=device)

    total_loss = self.cls_weight * cls_loss + self.iou_weight * iou_loss + self.dfl_weight * dfl_loss

    loss_dict = {
        "cls_loss": cls_loss.item(),
        "iou_loss": iou_loss.item(),
        "dfl_loss": dfl_loss.item(),
        "total_loss": total_loss.item(),
    }

    return total_loss, loss_dict

modern_yolonas.training.metrics.DetectionMetrics

COCO-style object-detection metrics backed by :class:torchmetrics.detection.MeanAveragePrecision.

Computes mAP @IoU[0.50:0.95], mAP@0.50, mAP@0.75, per-size mAP, and maximum average recall (mAR) for one, ten, and one-hundred max-detections per image — i.e. the same set reported by the COCO benchmark.

Parameters:

Name	Type	Description	Default
class_metrics	bool	When True, also produce per-class AP/AR in addition to the global averages.	False
extended_summary	bool	Forward the extended_summary flag to torchmetrics for extra diagnostic statistics.	False
device	str | device	Device on which internal accumulation tensors are kept. Use the training device to avoid cross-device copies during validation.	'cpu'

Example::

metrics = DetectionMetrics(device="cuda")
for images, targets in val_loader:
    preds = decode_predictions(model(images))   # list of dicts
    gts   = convert_targets(targets)            # list of dicts
    metrics.update(preds, gts)

scores = metrics.compute()
# {"mAP": 0.423, "mAP_50": 0.612, "mAP_75": 0.448, ...}
metrics.reset()

Source code in src/modern_yolonas/training/metrics.py

class DetectionMetrics:
    """COCO-style object-detection metrics backed by :class:`torchmetrics.detection.MeanAveragePrecision`.

    Computes mAP @IoU[0.50:0.95], mAP@0.50, mAP@0.75, per-size mAP, and maximum
    average recall (mAR) for one, ten, and one-hundred max-detections per image —
    i.e. the same set reported by the COCO benchmark.

    Args:
        class_metrics: When ``True``, also produce per-class AP/AR in addition to
            the global averages.
        extended_summary: Forward the ``extended_summary`` flag to torchmetrics for
            extra diagnostic statistics.
        device: Device on which internal accumulation tensors are kept.  Use the
            training device to avoid cross-device copies during validation.

    Example::

        metrics = DetectionMetrics(device="cuda")
        for images, targets in val_loader:
            preds = decode_predictions(model(images))   # list of dicts
            gts   = convert_targets(targets)            # list of dicts
            metrics.update(preds, gts)

        scores = metrics.compute()
        # {"mAP": 0.423, "mAP_50": 0.612, "mAP_75": 0.448, ...}
        metrics.reset()
    """

    def __init__(
        self,
        class_metrics: bool = False,
        extended_summary: bool = False,
        device: str | torch.device = "cpu",
    ) -> None:
        from torchmetrics.detection import MeanAveragePrecision

        self._metric = MeanAveragePrecision(
            box_format="xyxy",
            iou_type="bbox",
            class_metrics=class_metrics,
            extended_summary=extended_summary,
        ).to(torch.device(device))

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def update(self, preds: list[dict], targets: list[dict]) -> None:
        """Accumulate predictions and ground-truth for one batch.

        Args:
            preds: One dict per image, each containing:

                - ``"boxes"``: ``FloatTensor[N, 4]`` in *xyxy* pixel coords
                  (x1, y1, x2, y2).
                - ``"scores"``: ``FloatTensor[N]`` confidence values in ``[0, 1]``.
                - ``"labels"``: ``IntTensor[N]`` zero-based class indices.

            targets: One dict per image, each containing:

                - ``"boxes"``: ``FloatTensor[M, 4]`` in *xyxy* pixel coords.
                - ``"labels"``: ``IntTensor[M]`` zero-based class indices.
        """
        self._metric.update(preds, targets)

    def compute(self) -> dict[str, float]:
        """Compute metrics over all accumulated predictions and return a flat dict.

        Returns:
            A ``dict[str, float]`` with three keys:

            ``mAP``
                Mean AP averaged over IoU thresholds 0.50–0.95 (primary COCO metric).
            ``mAP_50``
                Mean AP at IoU = 0.50 (PASCAL VOC-style, a.k.a. the IoU metric).
            ``mAR_100``
                Max recall given at most 100 detections per image.
        """
        raw = self._metric.compute()

        def _scalar(key: str) -> float:
            v = raw.get(key, torch.tensor(0.0))
            return float(v.item() if isinstance(v, Tensor) else v)

        return {
            "mAP":    _scalar("map"),
            "mAP_50": _scalar("map_50"),
            "mAR_100": _scalar("mar_100"),
        }

    def reset(self) -> None:
        """Reset all accumulated state, ready for the next evaluation round."""
        self._metric.reset()

update(preds, targets)

Accumulate predictions and ground-truth for one batch.

Parameters:

Name	Type	Description	Default
preds	list[dict]	One dict per image, each containing: "boxes": FloatTensor[N, 4] in xyxy pixel coords (x1, y1, x2, y2). "scores": FloatTensor[N] confidence values in [0, 1]. "labels": IntTensor[N] zero-based class indices.	required
targets	list[dict]	One dict per image, each containing: "boxes": FloatTensor[M, 4] in xyxy pixel coords. "labels": IntTensor[M] zero-based class indices.	required

Source code in src/modern_yolonas/training/metrics.py

def update(self, preds: list[dict], targets: list[dict]) -> None:
    """Accumulate predictions and ground-truth for one batch.

    Args:
        preds: One dict per image, each containing:

            - ``"boxes"``: ``FloatTensor[N, 4]`` in *xyxy* pixel coords
              (x1, y1, x2, y2).
            - ``"scores"``: ``FloatTensor[N]`` confidence values in ``[0, 1]``.
            - ``"labels"``: ``IntTensor[N]`` zero-based class indices.

        targets: One dict per image, each containing:

            - ``"boxes"``: ``FloatTensor[M, 4]`` in *xyxy* pixel coords.
            - ``"labels"``: ``IntTensor[M]`` zero-based class indices.
    """
    self._metric.update(preds, targets)

compute()

Compute metrics over all accumulated predictions and return a flat dict.

Returns:

Type	Description
dict[str, float]	A dict[str, float] with three keys:
dict[str, float]	mAP Mean AP averaged over IoU thresholds 0.50–0.95 (primary COCO metric).
dict[str, float]	mAP_50 Mean AP at IoU = 0.50 (PASCAL VOC-style, a.k.a. the IoU metric).
dict[str, float]	mAR_100 Max recall given at most 100 detections per image.

Source code in src/modern_yolonas/training/metrics.py

def compute(self) -> dict[str, float]:
    """Compute metrics over all accumulated predictions and return a flat dict.

    Returns:
        A ``dict[str, float]`` with three keys:

        ``mAP``
            Mean AP averaged over IoU thresholds 0.50–0.95 (primary COCO metric).
        ``mAP_50``
            Mean AP at IoU = 0.50 (PASCAL VOC-style, a.k.a. the IoU metric).
        ``mAR_100``
            Max recall given at most 100 detections per image.
    """
    raw = self._metric.compute()

    def _scalar(key: str) -> float:
        v = raw.get(key, torch.tensor(0.0))
        return float(v.item() if isinstance(v, Tensor) else v)

    return {
        "mAP":    _scalar("map"),
        "mAP_50": _scalar("map_50"),
        "mAR_100": _scalar("mar_100"),
    }

reset()

Reset all accumulated state, ready for the next evaluation round.

Source code in src/modern_yolonas/training/metrics.py

def reset(self) -> None:
    """Reset all accumulated state, ready for the next evaluation round."""
    self._metric.reset()

Callbacks

modern_yolonas.training.callbacks.EMACallback

Bases: Callback

Exponential Moving Average of model weights.

Maintains a shadow copy of model parameters updated with an exponential moving average. During validation and checkpointing, EMA weights are swapped in.

Parameters:

Name	Type	Description	Default
decay	float	EMA decay factor.	0.9997
warmup_steps	int	Steps before reaching full decay.	2000

Source code in src/modern_yolonas/training/callbacks.py

class EMACallback(L.Callback):
    """Exponential Moving Average of model weights.

    Maintains a shadow copy of model parameters updated with an exponential
    moving average. During validation and checkpointing, EMA weights are
    swapped in.

    Args:
        decay: EMA decay factor.
        warmup_steps: Steps before reaching full decay.
    """

    def __init__(self, decay: float = 0.9997, warmup_steps: int = 2000):
        super().__init__()
        self.decay = decay
        self.warmup_steps = warmup_steps
        self.updates = 0
        self.ema_model: nn.Module | None = None
        self._orig_state: dict | None = None

    def on_fit_start(self, trainer: L.Trainer, pl_module: L.LightningModule):
        self.ema_model = deepcopy(pl_module.model).eval()
        for p in self.ema_model.parameters():
            p.requires_grad_(False)

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        self.updates += 1
        d = self.decay * (1 - math.exp(-self.updates / self.warmup_steps))
        model_sd = pl_module.model.state_dict()
        for k, v in self.ema_model.state_dict().items():
            if v.dtype.is_floating_point:
                v.mul_(d).add_(model_sd[k], alpha=1.0 - d)

    def _swap_ema_in(self, pl_module: L.LightningModule):
        self._orig_state = {k: v.clone() for k, v in pl_module.model.state_dict().items()}
        pl_module.model.load_state_dict(self.ema_model.state_dict())

    def _swap_ema_out(self, pl_module: L.LightningModule):
        if self._orig_state is not None:
            pl_module.model.load_state_dict(self._orig_state)
            self._orig_state = None

    def on_validation_start(self, trainer, pl_module):
        if self.ema_model is not None:
            self._swap_ema_in(pl_module)

    def on_validation_end(self, trainer, pl_module):
        self._swap_ema_out(pl_module)

    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
        if self.ema_model is not None:
            checkpoint["ema_state_dict"] = self.ema_model.state_dict()
            checkpoint["ema_updates"] = self.updates
            # Save EMA weights as the model weights in the checkpoint
            ema_sd = self.ema_model.state_dict()
            prefix = "model."
            for k, v in ema_sd.items():
                checkpoint["state_dict"][prefix + k] = v

    def on_load_checkpoint(self, trainer, pl_module, checkpoint):
        if "ema_state_dict" in checkpoint:
            self.updates = checkpoint.get("ema_updates", 0)
            # ema_model is created in on_fit_start; defer loading
            self._pending_ema_state = checkpoint["ema_state_dict"]

    def on_train_start(self, trainer, pl_module):
        if hasattr(self, "_pending_ema_state") and self._pending_ema_state is not None:
            if self.ema_model is not None:
                self.ema_model.load_state_dict(self._pending_ema_state)
            self._pending_ema_state = None

modern_yolonas.training.callbacks.CloseMosaicCallback

Bases: Callback

Disable Mosaic and Mixup augmentations in the final N epochs.

Accesses the train dataset's transforms pipeline and calls disable_mosaic_mixup(). Idempotent — safe to call every epoch.

Parameters:

Name	Type	Description	Default
close_mosaic_epochs	int	Number of final epochs without Mosaic/Mixup.	15

Source code in src/modern_yolonas/training/callbacks.py

class CloseMosaicCallback(L.Callback):
    """Disable Mosaic and Mixup augmentations in the final N epochs.

    Accesses the train dataset's transforms pipeline and calls
    ``disable_mosaic_mixup()``. Idempotent — safe to call every epoch.

    Args:
        close_mosaic_epochs: Number of final epochs without Mosaic/Mixup.
    """

    def __init__(self, close_mosaic_epochs: int = 15):
        super().__init__()
        self.close_mosaic_epochs = close_mosaic_epochs

    def on_train_epoch_start(self, trainer: L.Trainer, pl_module: L.LightningModule):
        epoch = trainer.current_epoch
        max_epochs = trainer.max_epochs
        if epoch >= max_epochs - self.close_mosaic_epochs:
            # Walk through dataloader to find the dataset transforms
            dataloader = trainer.train_dataloader
            if dataloader is None:
                return
            dataset = dataloader.dataset
            if hasattr(dataset, 'transforms') and hasattr(dataset.transforms, 'disable_mosaic_mixup'):
                dataset.transforms.disable_mosaic_mixup()

modern_yolonas.training.callbacks.QATCallback

Bases: Callback

Freeze BatchNorm and observers during Quantization-Aware Training.

Parameters:

Name	Type	Description	Default
freeze_bn_after_epoch	int	Freeze BN running stats after this epoch.	3
freeze_observer_after_epoch	int	Stop updating fake-quant ranges after this epoch.	5

Source code in src/modern_yolonas/training/callbacks.py

class QATCallback(L.Callback):
    """Freeze BatchNorm and observers during Quantization-Aware Training.

    Args:
        freeze_bn_after_epoch: Freeze BN running stats after this epoch.
        freeze_observer_after_epoch: Stop updating fake-quant ranges after this epoch.
    """

    def __init__(self, freeze_bn_after_epoch: int = 3, freeze_observer_after_epoch: int = 5):
        super().__init__()
        self.freeze_bn_after_epoch = freeze_bn_after_epoch
        self.freeze_observer_after_epoch = freeze_observer_after_epoch

    def on_train_epoch_start(self, trainer: L.Trainer, pl_module: L.LightningModule):
        epoch = trainer.current_epoch

        if epoch >= self.freeze_bn_after_epoch:
            for m in pl_module.model.backbone_neck.modules():
                if isinstance(m, nn.BatchNorm2d):
                    m.eval()

        if epoch >= self.freeze_observer_after_epoch:
            pl_module.model.backbone_neck.apply(torch.ao.quantization.disable_observer)

Checkpoints

modern_yolonas.weights.extract_model_state_dict(checkpoint_path, map_location='cpu')

Load model weights from either a Lightning .ckpt or legacy .pt checkpoint.

Handles: - Lightning format: state_dict keys prefixed with model. - Legacy format: model_state_dict key or EMA ema.ema_state_dict - Plain state_dict (e.g. from super-gradients pretrained weights)

Source code in src/modern_yolonas/weights.py

def extract_model_state_dict(checkpoint_path: str | Path, map_location: str = "cpu") -> dict:
    """Load model weights from either a Lightning .ckpt or legacy .pt checkpoint.

    Handles:
    - Lightning format: state_dict keys prefixed with ``model.``
    - Legacy format: ``model_state_dict`` key or EMA ``ema.ema_state_dict``
    - Plain state_dict (e.g. from super-gradients pretrained weights)
    """
    ckpt = torch.load(checkpoint_path, map_location=map_location, weights_only=False)

    # Lightning checkpoint format
    if "state_dict" in ckpt:
        sd = ckpt["state_dict"]
        # Strip 'model.' prefix added by LightningModule
        prefix = "model."
        stripped = {}
        for k, v in sd.items():
            if k.startswith(prefix):
                stripped[k[len(prefix):]] = v
            else:
                stripped[k] = v
        return stripped

    # Legacy format with EMA (prefer EMA weights if available)
    if "ema" in ckpt and "ema_state_dict" in ckpt["ema"]:
        return ckpt["ema"]["ema_state_dict"]

    # Legacy format with model_state_dict key
    if "model_state_dict" in ckpt:
        return ckpt["model_state_dict"]

    # Plain state_dict
    return ckpt