birder 0.2.1__tar.gz → 0.2.2__tar.gz

{birder-0.2.1 → birder-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder
-Version: 0.2.1
+Version: 0.2.2
 Summary: An open-source computer vision framework for wildlife image analysis, featuring state-of-the-art models for species classification and detection.
 Author: Ofer Hasson
 License-Expression: Apache-2.0
@@ -62,6 +62,7 @@ Requires-Dist: MonkeyType~=23.3.0; extra == "dev"
 Requires-Dist: mypy~=1.19.1; extra == "dev"
 Requires-Dist: parameterized~=0.9.0; extra == "dev"
 Requires-Dist: pylint~=4.0.4; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
 Requires-Dist: requests~=2.32.5; extra == "dev"
 Requires-Dist: safetensors~=0.7.0; extra == "dev"
 Requires-Dist: setuptools; extra == "dev"

birder-0.2.2/birder/adversarial/__init__.py

@@ -0,0 +1,13 @@
+from birder.adversarial.base import AttackResult
+from birder.adversarial.deepfool import DeepFool
+from birder.adversarial.fgsm import FGSM
+from birder.adversarial.pgd import PGD
+from birder.adversarial.simba import SimBA
+
+__all__ = [
+    "AttackResult",
+    "DeepFool",
+    "FGSM",
+    "PGD",
+    "SimBA",
+]

birder-0.2.2/birder/adversarial/base.py

@@ -0,0 +1,101 @@
+from dataclasses import dataclass
+from typing import Optional
+from typing import Protocol
+
+import torch
+
+from birder.data.transforms.classification import RGBType
+
+
+@dataclass(frozen=True)
+class AttackResult:
+    adv_inputs: torch.Tensor
+    adv_logits: torch.Tensor
+    perturbation: torch.Tensor
+    logits: Optional[torch.Tensor] = None
+    success: Optional[torch.Tensor] = None
+    num_queries: Optional[int] = None
+
+
+class Attack(Protocol):
+    def __call__(self, input_tensor: torch.Tensor, target: Optional[torch.Tensor]) -> AttackResult: ...
+
+
+def _to_channel_tensor(
+    values: tuple[float, float, float], device: Optional[torch.device], dtype: Optional[torch.dtype]
+) -> torch.Tensor:
+    return torch.tensor(values, device=device, dtype=dtype).view(1, -1, 1, 1)
+
+
+def normalized_bounds(
+    rgb_stats: RGBType, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    mean = _to_channel_tensor(rgb_stats["mean"], device=device, dtype=dtype)
+    std = _to_channel_tensor(rgb_stats["std"], device=device, dtype=dtype)
+    min_val = (0.0 - mean) / std
+    max_val = (1.0 - mean) / std
+
+    return (min_val, max_val)
+
+
+def pixel_eps_to_normalized(
+    eps: float | torch.Tensor,
+    rgb_stats: RGBType,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    eps_tensor = torch.as_tensor(eps, device=device, dtype=dtype)
+    std = _to_channel_tensor(rgb_stats["std"], device=eps_tensor.device, dtype=eps_tensor.dtype)
+
+    if eps_tensor.numel() == 1:
+        eps_tensor = eps_tensor.reshape(1, 1, 1, 1)
+    else:
+        eps_tensor = eps_tensor.reshape(1, -1, 1, 1)
+
+    return eps_tensor / std
+
+
+def clamp_normalized(inputs: torch.Tensor, rgb_stats: RGBType) -> torch.Tensor:
+    (min_val, max_val) = normalized_bounds(rgb_stats, device=inputs.device, dtype=inputs.dtype)
+    return torch.clamp(inputs, min=min_val, max=max_val)
+
+
+def predict_labels(logits: torch.Tensor) -> torch.Tensor:
+    return torch.argmax(logits, dim=1)
+
+
+def validate_target(
+    target: Optional[torch.Tensor], batch_size: int, num_classes: int, device: torch.device
+) -> Optional[torch.Tensor]:
+    if target is None:
+        return None
+
+    target = target.to(device=device, dtype=torch.long)
+    if target.ndim == 0:
+        target = target.view(1)
+
+    if target.shape[0] != batch_size:
+        raise ValueError(f"Target shape {target.shape[0]} must match batch size {batch_size}")
+
+    if torch.any(target < 0) or torch.any(target >= num_classes):
+        raise ValueError(f"Target values must be in range [0, {num_classes})")
+
+    return target
+
+
+def attack_success(
+    logits: torch.Tensor,
+    adv_logits: torch.Tensor,
+    targeted: bool,
+    target: Optional[torch.Tensor] = None,
+    labels: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    adv_pred = predict_labels(adv_logits)
+    if targeted is True:
+        if target is None:
+            raise ValueError("Target labels required for targeted attacks")
+
+        return adv_pred.eq(target)
+
+    base_labels = labels if labels is not None else predict_labels(logits)
+    return adv_pred.ne(base_labels)

birder-0.2.2/birder/adversarial/deepfool.py

@@ -0,0 +1,173 @@
+"""
+DeepFool
+
+Paper "DeepFool: a simple and accurate method to fool deep neural networks", https://arxiv.org/abs/1511.04599
+"""
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+from birder.adversarial.base import AttackResult
+from birder.adversarial.base import attack_success
+from birder.adversarial.base import clamp_normalized
+from birder.adversarial.base import predict_labels
+from birder.adversarial.base import validate_target
+from birder.data.transforms.classification import RGBType
+
+GRAD_EPS = 1e-12
+
+
+class DeepFool:
+    def __init__(
+        self, net: nn.Module, num_classes: int = 10, overshoot: float = 0.02, max_iter: int = 50, *, rgb_stats: RGBType
+    ) -> None:
+        if num_classes < 2:
+            raise ValueError("num_classes must be at least 2")
+        if max_iter <= 0:
+            raise ValueError("max_iter must be positive")
+        if overshoot < 0:
+            raise ValueError("overshoot must be non-negative")
+
+        self.net = net.eval()
+        self.num_classes = num_classes
+        self.overshoot = overshoot
+        self.max_iter = max_iter
+        self.rgb_stats = rgb_stats
+
+    def __call__(self, input_tensor: torch.Tensor, target: Optional[torch.Tensor]) -> AttackResult:
+        inputs = input_tensor.detach()
+        with torch.no_grad():
+            logits = self.net(inputs)
+
+        target_labels = (
+            validate_target(target, inputs.shape[0], logits.shape[1], inputs.device) if target is not None else None
+        )
+        targeted = target_labels is not None
+
+        adv_inputs_list = []
+        for idx in range(inputs.size(0)):
+            target_label = target_labels[idx : idx + 1] if target_labels is not None else None
+            adv_input = self._attack_single(inputs[idx : idx + 1], logits[idx : idx + 1], target_label)
+            adv_inputs_list.append(adv_input)
+
+        adv_inputs = torch.concat(adv_inputs_list, dim=0)
+        with torch.no_grad():
+            adv_logits = self.net(adv_inputs)
+
+        success = attack_success(
+            logits,
+            adv_logits,
+            targeted,
+            target=target_labels if targeted else None,
+        )
+
+        return AttackResult(
+            adv_inputs=adv_inputs,
+            adv_logits=adv_logits,
+            perturbation=adv_inputs - inputs,
+            logits=logits.detach(),
+            success=success,
+        )
+
+    def _attack_single(
+        self, inputs: torch.Tensor, logits: torch.Tensor, target_label: Optional[torch.Tensor]
+    ) -> torch.Tensor:
+        adv_inputs = inputs.clone()
+        original_label = int(predict_labels(logits).item())
+        targeted = target_label is not None
+        for _ in range(self.max_iter):
+            adv_inputs.requires_grad_(True)
+            outputs = self.net(adv_inputs)
+            current_label = int(predict_labels(outputs).item())
+
+            if targeted is True:
+                assert target_label is not None
+                target_value = int(target_label.item())
+                if current_label == target_value:
+                    break
+
+                perturbation = self._targeted_perturbation(adv_inputs, outputs, current_label, target_value)
+
+            else:
+                if current_label != original_label:
+                    break
+
+                perturbation = self._untargeted_perturbation(adv_inputs, outputs, current_label)
+
+            if perturbation is None:
+                break
+
+            # Overshoot helps ensure boundary crossing
+            adv_inputs = adv_inputs.detach() + (1.0 + self.overshoot) * perturbation
+            adv_inputs = clamp_normalized(adv_inputs, self.rgb_stats)
+
+        return adv_inputs.detach()
+
+    def _targeted_perturbation(
+        self, adv_inputs: torch.Tensor, outputs: torch.Tensor, current_label: int, target_label: int
+    ) -> Optional[torch.Tensor]:
+        self.net.zero_grad(set_to_none=True)
+        grad_current = torch.autograd.grad(outputs[0, current_label], adv_inputs, retain_graph=True)[0]
+        grad_target = torch.autograd.grad(outputs[0, target_label], adv_inputs, retain_graph=False)[0]
+
+        # Direction toward the target boundary
+        w = grad_target - grad_current
+        w_norm = torch.norm(w.view(-1))
+        if w_norm.item() < GRAD_EPS:
+            return None
+
+        # Distance to the decision boundary
+        f = outputs[0, target_label] - outputs[0, current_label]
+        perturbation = (f.abs() / (w_norm**2 + GRAD_EPS)) * w
+
+        return perturbation
+
+    def _untargeted_perturbation(
+        self, adv_inputs: torch.Tensor, outputs: torch.Tensor, current_label: int
+    ) -> Optional[torch.Tensor]:
+        # Search the top-k competing classes
+        top_k = min(self.num_classes, outputs.shape[1])
+        top_indices = torch.topk(outputs, k=top_k, dim=1).indices[0]
+        candidate_labels = [int(idx) for idx in top_indices if int(idx) != current_label]
+
+        if len(candidate_labels) == 0:
+            return None
+
+        self.net.zero_grad(set_to_none=True)
+        grad_current = torch.autograd.grad(outputs[0, current_label], adv_inputs, retain_graph=True)[0]
+
+        # Track the closest decision boundary
+        best_dist = None
+        best_w = None
+        best_f = None
+        for idx, label in enumerate(candidate_labels):
+            # Keep the graph until the last class
+            retain_graph = idx != len(candidate_labels) - 1
+            grad_other = torch.autograd.grad(outputs[0, label], adv_inputs, retain_graph=retain_graph)[0]
+
+            w_k = grad_other - grad_current
+            w_norm = torch.norm(w_k.view(-1))
+            if w_norm.item() < GRAD_EPS:
+                continue
+
+            f_k = outputs[0, label] - outputs[0, current_label]
+            dist = f_k.abs() / (w_norm + GRAD_EPS)
+
+            if best_dist is None or dist < best_dist:
+                best_dist = dist
+                best_w = w_k
+                best_f = f_k
+
+        if best_w is None or best_f is None:
+            return None
+
+        # Minimal perturbation toward the closest boundary
+        best_w_norm = torch.norm(best_w.view(-1))
+        if best_w_norm.item() < GRAD_EPS:
+            return None
+
+        perturbation = (best_f.abs() / (best_w_norm**2 + GRAD_EPS)) * best_w
+
+        return perturbation

birder-0.2.2/birder/adversarial/fgsm.py

@@ -0,0 +1,67 @@
+"""
+Fast Gradient Sign Method (FGSM)
+
+Paper "Explaining and Harnessing Adversarial Examples", https://arxiv.org/abs/1412.6572
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from birder.adversarial.base import AttackResult
+from birder.adversarial.base import attack_success
+from birder.adversarial.base import clamp_normalized
+from birder.adversarial.base import pixel_eps_to_normalized
+from birder.adversarial.base import predict_labels
+from birder.adversarial.base import validate_target
+from birder.data.transforms.classification import RGBType
+
+
+class FGSM:
+    def __init__(self, net: nn.Module, eps: float, *, rgb_stats: RGBType) -> None:
+        self.net = net.eval()
+        self.eps = eps
+        self.rgb_stats = rgb_stats
+
+    def __call__(self, input_tensor: torch.Tensor, target: Optional[torch.Tensor]) -> AttackResult:
+        inputs = input_tensor.detach().clone()
+        inputs.requires_grad_(True)
+
+        logits = self.net(inputs)
+        targeted = target is not None
+        if targeted is True:
+            target = validate_target(target, inputs.shape[0], logits.shape[1], inputs.device)
+        else:
+            target = predict_labels(logits)
+
+        loss = F.cross_entropy(logits, target)
+        (grad,) = torch.autograd.grad(loss, inputs, retain_graph=False, create_graph=False)
+        eps_norm = pixel_eps_to_normalized(self.eps, self.rgb_stats, device=inputs.device, dtype=inputs.dtype)
+
+        # Targeted steps descend toward target, untargeted ascend away from original
+        if targeted is True:
+            direction = -1.0
+        else:
+            direction = 1.0
+
+        perturbation = direction * eps_norm * grad.sign()
+        adv_inputs = clamp_normalized(inputs + perturbation, self.rgb_stats)
+        with torch.no_grad():
+            adv_logits = self.net(adv_inputs)
+
+        success = attack_success(
+            logits.detach(),
+            adv_logits,
+            targeted,
+            target=target if targeted else None,
+        )
+
+        return AttackResult(
+            adv_inputs=adv_inputs,
+            adv_logits=adv_logits,
+            perturbation=adv_inputs - inputs,
+            logits=logits.detach(),
+            success=success,
+        )

birder-0.2.2/birder/adversarial/pgd.py

@@ -0,0 +1,105 @@
+"""
+Projected Gradient Descent (PGD)
+
+Paper "Towards Deep Learning Models Resistant to Adversarial Attacks", https://arxiv.org/abs/1706.06083
+"""
+
+# Reference license: MIT
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from birder.adversarial.base import AttackResult
+from birder.adversarial.base import attack_success
+from birder.adversarial.base import clamp_normalized
+from birder.adversarial.base import pixel_eps_to_normalized
+from birder.adversarial.base import predict_labels
+from birder.adversarial.base import validate_target
+from birder.data.transforms.classification import RGBType
+
+
+class PGD:
+    def __init__(
+        self,
+        net: nn.Module,
+        eps: float,
+        steps: int = 10,
+        step_size: Optional[float] = None,
+        random_start: bool = False,
+        *,
+        rgb_stats: RGBType,
+    ) -> None:
+        if steps <= 0:
+            raise ValueError("steps must be a positive integer")
+
+        self.net = net.eval()
+        self.eps = eps
+        self.steps = steps
+        if step_size is not None:
+            self.step_size = step_size
+        else:
+            self.step_size = eps / steps
+
+        self.random_start = random_start
+        self.rgb_stats = rgb_stats
+
+        if self.step_size <= 0:
+            raise ValueError("step_size must be positive")
+
+    def __call__(self, input_tensor: torch.Tensor, target: Optional[torch.Tensor]) -> AttackResult:
+        inputs = input_tensor.detach()
+        with torch.no_grad():
+            logits = self.net(inputs)
+
+        targeted = target is not None
+        if targeted:
+            target = validate_target(target, inputs.shape[0], logits.shape[1], inputs.device)
+        else:
+            target = predict_labels(logits)
+
+        eps_norm = pixel_eps_to_normalized(self.eps, self.rgb_stats, device=inputs.device, dtype=inputs.dtype)
+        step_norm = pixel_eps_to_normalized(self.step_size, self.rgb_stats, device=inputs.device, dtype=inputs.dtype)
+
+        # Targeted steps descend toward target, untargeted ascend away from original
+        if targeted is True:
+            direction = -1.0
+        else:
+            direction = 1.0
+
+        adv_inputs = inputs.clone()
+        if self.random_start is True:
+            # Random start inside the epsilon ball
+            adv_inputs = adv_inputs + torch.empty_like(adv_inputs).uniform_(-1.0, 1.0) * eps_norm
+            adv_inputs = clamp_normalized(adv_inputs, self.rgb_stats)
+
+        for _ in range(self.steps):
+            adv_inputs.requires_grad_(True)
+            adv_logits = self.net(adv_inputs)
+            loss = F.cross_entropy(adv_logits, target)
+            (grad,) = torch.autograd.grad(loss, adv_inputs, retain_graph=False, create_graph=False)
+            adv_inputs = adv_inputs.detach() + direction * step_norm * grad.sign()
+
+            # Project back into the epsilon ball around the original input.
+            delta = torch.clamp(adv_inputs - inputs, min=-eps_norm, max=eps_norm)
+            adv_inputs = clamp_normalized(inputs + delta, self.rgb_stats)
+
+        with torch.no_grad():
+            adv_logits = self.net(adv_inputs)
+
+        success = attack_success(
+            logits.detach(),
+            adv_logits,
+            targeted,
+            target=target if targeted else None,
+        )
+
+        return AttackResult(
+            adv_inputs=adv_inputs,
+            adv_logits=adv_logits,
+            perturbation=adv_inputs - inputs,
+            logits=logits.detach(),
+            success=success,
+        )

birder-0.2.2/birder/adversarial/simba.py

@@ -0,0 +1,172 @@
+"""
+SimBA (Simple Black-box Attack)
+
+Paper "Simple Black-box Adversarial Attacks", https://arxiv.org/abs/1905.07121
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from birder.adversarial.base import AttackResult
+from birder.adversarial.base import attack_success
+from birder.adversarial.base import clamp_normalized
+from birder.adversarial.base import pixel_eps_to_normalized
+from birder.adversarial.base import predict_labels
+from birder.adversarial.base import validate_target
+from birder.data.transforms.classification import RGBType
+
+
+class SimBA:
+    def __init__(self, net: nn.Module, step_size: float, max_iter: int = 1000, *, rgb_stats: RGBType) -> None:
+        if step_size <= 0:
+            raise ValueError("step_size must be positive")
+        if max_iter <= 0:
+            raise ValueError("max_iter must be positive")
+
+        self.net = net.eval()
+        self.step_size = step_size
+        self.max_iter = max_iter
+        self.rgb_stats = rgb_stats
+
+    def __call__(self, input_tensor: torch.Tensor, target: Optional[torch.Tensor]) -> AttackResult:
+        inputs = input_tensor.detach()
+        with torch.no_grad():
+            logits = self.net(inputs)
+
+        labels = predict_labels(logits)
+        target_labels = (
+            validate_target(target, inputs.shape[0], logits.shape[1], inputs.device) if target is not None else None
+        )
+        targeted = target_labels is not None
+
+        adv_inputs_list = []
+        total_queries = 0
+        for idx in range(inputs.size(0)):
+            label = labels[idx : idx + 1]
+            target_label = target_labels[idx : idx + 1] if target_labels is not None else None
+            adv_input, num_queries = self._attack_single(inputs[idx : idx + 1], label, target_label)
+            adv_inputs_list.append(adv_input)
+            total_queries += num_queries
+
+        adv_inputs = torch.concat(adv_inputs_list, dim=0)
+        with torch.no_grad():
+            adv_logits = self.net(adv_inputs)
+
+        success = attack_success(
+            logits,
+            adv_logits,
+            targeted,
+            target=target_labels if targeted else None,
+        )
+
+        return AttackResult(
+            adv_inputs=adv_inputs,
+            adv_logits=adv_logits,
+            perturbation=adv_inputs - inputs,
+            logits=logits.detach(),
+            success=success,
+            num_queries=total_queries,
+        )
+
+    # pylint: disable=too-many-locals
+    def _attack_single(
+        self, inputs: torch.Tensor, label: torch.Tensor, target_label: Optional[torch.Tensor]
+    ) -> tuple[torch.Tensor, int]:
+        adv_inputs = inputs.clone()
+        num_queries = 1  # Baseline forward pass
+
+        with torch.no_grad():
+            current_logits = self.net(adv_inputs)
+            current_objective = self._compute_objective(current_logits, label, target_label)
+
+        if self._is_successful(current_logits, label, target_label):
+            return adv_inputs.detach(), num_queries
+
+        (_, channels, height, width) = adv_inputs.shape
+        num_dims = channels * height * width
+        step = pixel_eps_to_normalized(self.step_size, self.rgb_stats, device=adv_inputs.device, dtype=adv_inputs.dtype)
+        step_vals = step.view(-1)  # Per-channel steps
+        stride = height * width
+
+        perm = torch.randperm(num_dims, device=adv_inputs.device)
+        num_steps = min(self.max_iter, num_dims)
+
+        # Coordinate-wise search in random order
+        for flat_idx in perm[:num_steps]:
+            (c, rem) = divmod(int(flat_idx.item()), stride)
+            (h, w) = divmod(rem, width)
+            step_val = step_vals[c]
+
+            (candidate_inputs, candidate_logits, candidate_objective) = self._best_candidate(
+                adv_inputs, c, h, w, step_val, label, target_label
+            )
+            num_queries += 2
+
+            if candidate_objective < current_objective:
+                adv_inputs = candidate_inputs
+                current_logits = candidate_logits
+                current_objective = candidate_objective
+
+                if self._is_successful(current_logits, label, target_label) is True:
+                    break
+
+        return adv_inputs.detach(), num_queries
+
+    def _perturb_pixel(
+        self, inputs: torch.Tensor, channel: int, row: int, col: int, step: torch.Tensor
+    ) -> torch.Tensor:
+        adv_inputs = inputs.clone()
+        adv_inputs[0, channel, row, col] = adv_inputs[0, channel, row, col] + step
+        return clamp_normalized(adv_inputs, self.rgb_stats)
+
+    def _evaluate_candidate(
+        self, inputs: torch.Tensor, label: torch.Tensor, target_label: Optional[torch.Tensor]
+    ) -> tuple[torch.Tensor, float]:
+        with torch.no_grad():
+            logits = self.net(inputs)
+
+        return logits, self._compute_objective(logits, label, target_label)
+
+    def _best_candidate(
+        self,
+        inputs: torch.Tensor,
+        channel: int,
+        row: int,
+        col: int,
+        step: torch.Tensor,
+        label: torch.Tensor,
+        target_label: Optional[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, float]:
+        adv_plus = self._perturb_pixel(inputs, channel, row, col, step)
+        logits_plus, objective_plus = self._evaluate_candidate(adv_plus, label, target_label)
+
+        adv_minus = self._perturb_pixel(inputs, channel, row, col, -step)
+        logits_minus, objective_minus = self._evaluate_candidate(adv_minus, label, target_label)
+
+        if objective_plus <= objective_minus:
+            return adv_plus, logits_plus, objective_plus
+
+        return adv_minus, logits_minus, objective_minus
+
+    @staticmethod
+    def _compute_objective(
+        logits: torch.Tensor, original_label: torch.Tensor, target_label: Optional[torch.Tensor]
+    ) -> float:
+        # Lower objective is better in both modes
+        if target_label is not None:
+            return float(F.cross_entropy(logits, target_label).item())
+
+        return -float(F.cross_entropy(logits, original_label).item())
+
+    @staticmethod
+    def _is_successful(
+        logits: torch.Tensor, original_label: torch.Tensor, target_label: Optional[torch.Tensor]
+    ) -> bool:
+        pred = predict_labels(logits)
+        if target_label is not None:
+            return bool(pred.eq(target_label).item())
+
+        return bool(pred.ne(original_label).item())

{birder-0.2.1 → birder-0.2.2}/birder/common/training_cli.py

@@ -110,10 +110,13 @@ def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
         type=int,
         default=40,
         metavar="N",
-        help="decrease lr every step-size epochs (for step scheduler only)",
+        help="decrease lr every N epochs/steps (relative to after warmup, step scheduler only)",
     )
     group.add_argument(
-        "--lr-steps", type=int, nargs="+", help="decrease lr every step-size epochs (multistep scheduler only)"
+        "--lr-steps",
+        type=int,
+        nargs="+",
+        help="absolute epoch/step milestones when to decrease lr (multistep scheduler only)",
     )
     group.add_argument(
         "--lr-step-gamma",
@@ -391,7 +394,7 @@ def add_ema_args(
         "--model-ema-warmup",
         type=int,
         metavar="N",
-        help="number of epochs before EMA is applied (defaults to warmup epochs/iters, pass 0 to disable warmup)",
+        help="number of epochs/steps before EMA is applied (defaults to warmup epochs/steps, pass 0 to disable warmup)",
     )
 
 
@@ -656,6 +659,11 @@ def common_args_validation(args: argparse.Namespace) -> None:
                 f"but it is set to '{args.lr_scheduler_update}'"
             )
 
+    # EMA
+    if hasattr(args, "model_ema_steps") is True:
+        if args.model_ema_steps < 1:
+            raise ValidationError("--model-ema-steps must be >= 1")
+
     # Compile args, argument dependant
     if hasattr(args, "compile_teacher") is True:
         if args.compile is True and args.compile_teacher is True: