YoloNAS Model

modern_yolonas.model.YoloNAS

Bases: Module

YOLO-NAS object detection model.

Attributes backbone, neck, heads mirror the super-gradients top-level structure for state_dict key compatibility.

Source code in src/modern_yolonas/model.py

class YoloNAS(nn.Module):
    """YOLO-NAS object detection model.

    Attributes ``backbone``, ``neck``, ``heads`` mirror the super-gradients
    top-level structure for state_dict key compatibility.
    """

    def __init__(
        self,
        backbone: YoloNASBackbone,
        neck: YoloNASPANNeckWithC2,
        heads: NDFLHeads,
    ):
        super().__init__()
        self.backbone = backbone
        self.neck = neck
        self.heads = heads
        self._init_bn()

    def _init_bn(self, eps: float = 1e-3, momentum: float = 0.03):
        """Apply super-gradients-compatible BatchNorm settings globally."""
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eps = eps
                m.momentum = momentum

    def forward(self, x: Tensor):
        features = self.backbone(x)
        p3, p4, p5 = self.neck(features)
        return self.heads((p3, p4, p5))

    def forward_features(self, x: Tensor) -> dict[str, Tensor]:
        """Run backbone and neck and return every intermediate feature map.

        The detection head is skipped entirely — this is the representation the
        model builds *before* it commits to boxes and classes, which is what an
        image-retrieval or re-identification pipeline wants. See
        :class:`~modern_yolonas.inference.embed.YoloNASEmbedder` for the
        image-in, vector-out wrapper around it.

        Deliberately a separate method rather than a flag on ``forward``: the
        signature of ``forward`` is what ONNX export, the Frigate graph, the
        parity tests and the FX-traced quantization path all depend on.

        Args:
            x: Normalized input batch, ``[B, 3, H, W]``.

        Returns:
            ``{"c2","c3","c4","c5","p3","p4","p5"}`` → feature maps. ``c2``–``c5``
            are the backbone outputs (``c5`` is post-SPP); ``p3``–``p5`` are the
            neck's fused pyramid levels. At 640x640 the strides are 4, 8, 16, 32
            for ``c2``–``c5`` and 8, 16, 32 for ``p3``–``p5``.
        """
        c2, c3, c4, c5 = self.backbone(x)
        p3, p4, p5 = self.neck([c2, c3, c4, c5])
        return {"c2": c2, "c3": c3, "c4": c4, "c5": c5, "p3": p3, "p4": p4, "p5": p5}

    @classmethod
    def from_config(cls, variant: str, num_classes: int = 80) -> "YoloNAS":
        cfg = CONFIGS[variant]

        backbone = YoloNASBackbone(
            in_channels=3,
            stem_out_channels=cfg["backbone"]["stem_out_channels"],
            stages_config=cfg["backbone"]["stages"],
            spp_output_channels=cfg["backbone"]["spp_output_channels"],
            spp_k=cfg["backbone"]["spp_k"],
        )

        neck = YoloNASPANNeckWithC2(
            in_channels=backbone.out_channels,
            **cfg["neck"],
        )

        heads = NDFLHeads(
            num_classes=num_classes,
            in_channels=tuple(neck.out_channels),
            **cfg["heads"],
        )

        return cls(backbone=backbone, neck=neck, heads=heads)

forward_features(x)

Run backbone and neck and return every intermediate feature map.

The detection head is skipped entirely — this is the representation the model builds before it commits to boxes and classes, which is what an image-retrieval or re-identification pipeline wants. See :class:~modern_yolonas.inference.embed.YoloNASEmbedder for the image-in, vector-out wrapper around it.

Deliberately a separate method rather than a flag on forward: the signature of forward is what ONNX export, the Frigate graph, the parity tests and the FX-traced quantization path all depend on.

Parameters:

Name	Type	Description	Default
x	Tensor	Normalized input batch, [B, 3, H, W].	required

Returns:

Type	Description
dict[str, Tensor]	{"c2","c3","c4","c5","p3","p4","p5"} → feature maps. c2–c5
dict[str, Tensor]	are the backbone outputs (c5 is post-SPP); p3–p5 are the
dict[str, Tensor]	neck's fused pyramid levels. At 640x640 the strides are 4, 8, 16, 32
dict[str, Tensor]	for c2–c5 and 8, 16, 32 for p3–p5.

Source code in src/modern_yolonas/model.py

def forward_features(self, x: Tensor) -> dict[str, Tensor]:
    """Run backbone and neck and return every intermediate feature map.

    The detection head is skipped entirely — this is the representation the
    model builds *before* it commits to boxes and classes, which is what an
    image-retrieval or re-identification pipeline wants. See
    :class:`~modern_yolonas.inference.embed.YoloNASEmbedder` for the
    image-in, vector-out wrapper around it.

    Deliberately a separate method rather than a flag on ``forward``: the
    signature of ``forward`` is what ONNX export, the Frigate graph, the
    parity tests and the FX-traced quantization path all depend on.

    Args:
        x: Normalized input batch, ``[B, 3, H, W]``.

    Returns:
        ``{"c2","c3","c4","c5","p3","p4","p5"}`` → feature maps. ``c2``–``c5``
        are the backbone outputs (``c5`` is post-SPP); ``p3``–``p5`` are the
        neck's fused pyramid levels. At 640x640 the strides are 4, 8, 16, 32
        for ``c2``–``c5`` and 8, 16, 32 for ``p3``–``p5``.
    """
    c2, c3, c4, c5 = self.backbone(x)
    p3, p4, p5 = self.neck([c2, c3, c4, c5])
    return {"c2": c2, "c3": c3, "c4": c4, "c5": c5, "p3": p3, "p4": p4, "p5": p5}

Factory functions

modern_yolonas.yolo_nas_s(pretrained=False, num_classes=80)

Create a YOLO-NAS S model.

Source code in src/modern_yolonas/__init__.py

def yolo_nas_s(pretrained: bool = False, num_classes: int = 80) -> YoloNAS:
    """Create a YOLO-NAS S model."""
    model = YoloNAS.from_config("yolo_nas_s", num_classes=num_classes)
    if pretrained:
        load_pretrained(model, "yolo_nas_s", strict=(num_classes == 80))
    return model

modern_yolonas.yolo_nas_m(pretrained=False, num_classes=80)

Create a YOLO-NAS M model.

Source code in src/modern_yolonas/__init__.py

def yolo_nas_m(pretrained: bool = False, num_classes: int = 80) -> YoloNAS:
    """Create a YOLO-NAS M model."""
    model = YoloNAS.from_config("yolo_nas_m", num_classes=num_classes)
    if pretrained:
        load_pretrained(model, "yolo_nas_m", strict=(num_classes == 80))
    return model

modern_yolonas.yolo_nas_l(pretrained=False, num_classes=80)

Create a YOLO-NAS L model.

Source code in src/modern_yolonas/__init__.py

def yolo_nas_l(pretrained: bool = False, num_classes: int = 80) -> YoloNAS:
    """Create a YOLO-NAS L model."""
    model = YoloNAS.from_config("yolo_nas_l", num_classes=num_classes)
    if pretrained:
        load_pretrained(model, "yolo_nas_l", strict=(num_classes == 80))
    return model