Training API Reference

oriented_det.train

Training engine and utilities for oriented detection.

AugmentationConfig dataclass

Augmentation configuration. Keys ordered: limits first, then p_* (probabilities).

Source code in oriented_det/train/config.py

@dataclass
class AugmentationConfig:
    """Augmentation configuration. Keys ordered: limits first, then p_* (probabilities)."""
    brightness_limit: float = 0.2
    contrast_limit: float = 0.2
    gamma_limit: tuple[int, int] = (80, 120)
    gauss_noise_var_limit: tuple[float, float] = (10.0, 50.0)
    blur_limit: int = 3
    clahe_clip_limit: float = 4.0
    p_brightness_contrast: float = 0.5
    p_gamma: float = 0.3
    p_noise: float = 0.2
    p_blur: float = 0.2
    p_clahe: float = 0.3

CheckpointConfig dataclass

Checkpoint loading and best-model selection configuration.

Source code in oriented_det/train/config.py

@dataclass
class CheckpointConfig:
    """Checkpoint loading and best-model selection configuration."""
    load_from_checkpoint: Optional[Path] = None
    load_from_experiment: Optional[Path] = None
    # When True and load_from_experiment is still unset: use newest other run under runs/<model_type>/.
    # When False: never auto-fill experiment dir (null checkpoint + null experiment = no weights loaded).
    discover_previous_run: bool = False
    # True: prefer latest checkpoint_epoch_*.pth, restore optimizer/scheduler per load_* flags, continue epochs.
    # False: prefer best_* checkpoint, typically epoch 0 with fresh optimizer unless config says otherwise.
    resume_from_checkpoint_epoch: bool = True
    load_optimizer_state: bool = True
    # When True (default), restore LR scheduler from checkpoint on resume. Set False to rebuild
    # cosine rebuild on resume when load_scheduler_state=false (see tools/train.py).
    load_scheduler_state: bool = True
    load_include_prefixes: Optional[List[str]] = None  # Load only keys whose names start with these prefixes
    load_exclude_prefixes: Optional[List[str]] = None  # Drop keys whose names start with these prefixes
    start_epoch: int = 0
    # Best checkpoint: metric name used to decide which checkpoint to keep (e.g. "mAP" or "total_loss")
    best_metric: Optional[str] = None  # Default in train: "total_loss"; use "mAP" to save at best validation mAP
    higher_is_better: Optional[bool] = None  # True for mAP, False for total_loss; if None, inferred from best_metric

CheckpointManager

Manages model checkpointing with robust error handling.

Source code in oriented_det/train/utils.py

class CheckpointManager:
    """Manages model checkpointing with robust error handling."""

    def __init__(
        self,
        checkpoint_dir: str | Path,
        *,
        keep_last_n: int = 5,
        best_metric: Optional[str] = None,
        higher_is_better: bool = True,
    ):
        _require_torch()
        self.checkpoint_dir = Path(checkpoint_dir)
        # Directory is created on first save() so multi-GPU training creates it only on rank 0
        self.keep_last_n = keep_last_n
        self.best_metric = best_metric
        self.higher_is_better = higher_is_better
        self.best_value: Optional[float] = None
        self.best_checkpoint_path: Optional[Path] = None
        self.checkpoint_history: List[Path] = []

    def save(
        self,
        model: nn.Module,
        optimizer: Optional[torch.optim.Optimizer] = None,
        scheduler: Optional[Any] = None,
        epoch: Optional[int] = None,
        metrics: Optional[Dict[str, float]] = None,
        *,
        suffix: str = "",
        custom_filename: Optional[str] = None,
    ) -> Path:
        """Save a checkpoint.

        If custom_filename is set (e.g. "best_mAP_0.42.pth"), that name is used
        and the checkpoint is not added to history (treated as a best checkpoint).

        Returns path to saved checkpoint.
        """
        _require_torch()
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)

        checkpoint = {
            "model_state_dict": model.state_dict(),
            "epoch": epoch,
            "metrics": metrics or {},
        }

        if optimizer is not None:
            checkpoint["optimizer_state_dict"] = optimizer.state_dict()
        if scheduler is not None and hasattr(scheduler, "state_dict"):
            checkpoint["scheduler_state_dict"] = scheduler.state_dict()

        if custom_filename:
            filename = custom_filename if custom_filename.endswith(".pth") else f"{custom_filename}.pth"
        elif suffix:
            filename = f"checkpoint_{suffix}.pth"
        elif epoch is not None:
            filename = f"checkpoint_epoch_{epoch:04d}.pth"
        else:
            filename = f"checkpoint_{int(time.time())}.pth"

        checkpoint_path = self.checkpoint_dir / filename

        # Atomic write: write to temp file first, then rename. Retry on transient I/O errors
        # (e.g. NFS "unexpected pos" / "file write failed"). Fallback: write to local /tmp then move.
        temp_path = checkpoint_path.with_suffix(".tmp")
        max_attempts = 3
        last_error = None
        for attempt in range(max_attempts):
            try:
                torch.save(checkpoint, temp_path)
                temp_path.replace(checkpoint_path)
                last_error = None
                break
            except Exception as e:
                last_error = e
                if temp_path.exists():
                    try:
                        temp_path.unlink()
                    except OSError:
                        pass
                if attempt < max_attempts - 1:
                    time.sleep(1.0 * (attempt + 1))
                    continue
        if last_error is not None:
            # Fallback: write to local temp dir (avoids NFS/network fs write failures)
            local_tmp = Path(os.environ.get("TMPDIR", "/tmp"))
            if local_tmp.exists() and os.access(local_tmp, os.W_OK):
                fallback_temp = local_tmp / f"oriented_det_checkpoint_{os.getpid()}_{filename}.tmp"
                try:
                    torch.save(checkpoint, fallback_temp)
                    shutil.move(str(fallback_temp), str(checkpoint_path))
                except Exception as e2:
                    if fallback_temp.exists():
                        try:
                            fallback_temp.unlink()
                        except OSError:
                            pass
                    raise RuntimeError(
                        f"Failed to save checkpoint after {max_attempts} attempts and local fallback: {last_error!s}; fallback: {e2!s}"
                    ) from e2
            else:
                raise RuntimeError(f"Failed to save checkpoint: {last_error}") from last_error

        # Only add to history if it's not a best checkpoint (best is managed separately)
        if not custom_filename and suffix != "best":
            self.checkpoint_history.append(checkpoint_path)

            # Clean up old checkpoints (excluding best checkpoint)
            while len(self.checkpoint_history) > self.keep_last_n:
                old_checkpoint = self.checkpoint_history.pop(0)
                # Don't delete if it's the best checkpoint
                if old_checkpoint != self.best_checkpoint_path and old_checkpoint.exists():
                    old_checkpoint.unlink()

        return checkpoint_path

    def save_best(
        self,
        model: nn.Module,
        optimizer: Optional[torch.optim.Optimizer] = None,
        scheduler: Optional[Any] = None,
        epoch: Optional[int] = None,
        metrics: Optional[Dict[str, float]] = None,
    ) -> Optional[Path]:
        """Save checkpoint if it's the best according to best_metric.

        The best checkpoint is stored independently and is not subject to
        the keep_last_n cleanup, ensuring it's always preserved.
        """
        if self.best_metric is None or metrics is None:
            return None

        current_value = metrics.get(self.best_metric)
        if current_value is None:
            return None
        # Do not save best checkpoint when mAP was not computed (sentinel -1)
        if str(self.best_metric).strip().lower() == "map" and current_value == -1.0:
            return None

        is_better = (
            (self.best_value is None) or
            (self.higher_is_better and current_value > self.best_value) or
            (not self.higher_is_better and current_value < self.best_value)
        )

        if is_better:
            # Delete old best checkpoint if it exists
            if self.best_checkpoint_path is not None and self.best_checkpoint_path.exists():
                try:
                    self.best_checkpoint_path.unlink()
                except Exception as e:
                    # Log but don't fail if deletion fails
                    print(f"Warning: Failed to delete old best checkpoint {self.best_checkpoint_path}: {e}")

            self.best_value = current_value
            # Build filename with metric name and value, e.g. best_mAP_0.42.pth or best_total_loss_1.23.pth
            safe_metric = "".join(c if c.isalnum() or c == "_" else "_" for c in str(self.best_metric))
            if isinstance(current_value, float):
                value_str = f"{current_value:.2f}" if safe_metric.lower() == "map" else f"{current_value:.4f}"
            else:
                value_str = str(current_value)
            custom_filename = f"best_{safe_metric}_{value_str}.pth"
            self.best_checkpoint_path = self.save(
                model, optimizer, scheduler, epoch, metrics, custom_filename=custom_filename
            )
            return self.best_checkpoint_path

        return None

    def load(
        self,
        checkpoint_path: str | Path,
        model: nn.Module,
        optimizer: Optional[torch.optim.Optimizer] = None,
        scheduler: Optional[Any] = None,
        *,
        strict: bool = False,
        include_prefixes: Optional[List[str]] = None,
        exclude_prefixes: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """Load a checkpoint.

        Args:
            checkpoint_path: Path to checkpoint file
            model: Model to load weights into
            optimizer: Optional optimizer to load state into
            strict: If True, raise error on mismatch. If False, warn and load matching weights.

        Returns:
            Loaded checkpoint dict.
        """
        _require_torch()

        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        try:
            checkpoint = torch.load(checkpoint_path, map_location="cpu")
        except Exception as e:
            raise RuntimeError(f"Failed to load checkpoint: {e}") from e

        # Load model state dict
        try:
            # Support multiple checkpoint formats:
            # - this repo: {"model_state_dict": ...}
            # - MMDetection/MMRotate: {"state_dict": ...}
            # - raw state_dict: {param_name: tensor, ...}
            if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
                checkpoint_state_dict = checkpoint["model_state_dict"]
            elif isinstance(checkpoint, dict) and "state_dict" in checkpoint:
                checkpoint_state_dict = checkpoint["state_dict"]
            elif isinstance(checkpoint, dict) and checkpoint and all(isinstance(k, str) for k in checkpoint.keys()):
                checkpoint_state_dict = checkpoint
            else:
                raise KeyError("Unsupported checkpoint format: expected model_state_dict or state_dict.")

            # MMRotate → torchvision BackboneWithFPN naming bridge:
            # MMRotate backbone keys are typically "backbone.*" (ResNet), while this repo's
            # torchvision backbone lives under "backbone.body.*".
            model_state_keys = list(model.state_dict().keys())
            model_has_backbone_body = any(k.startswith("backbone.body.") for k in model_state_keys)
            ckpt_has_backbone_body = any(k.startswith("backbone.body.") for k in checkpoint_state_dict.keys())
            ckpt_has_backbone = any(k.startswith("backbone.") for k in checkpoint_state_dict.keys())
            if model_has_backbone_body and ckpt_has_backbone and not ckpt_has_backbone_body:
                remapped = {}
                for k, v in checkpoint_state_dict.items():
                    if k.startswith("backbone."):
                        suffix = k[len("backbone."):]
                        remapped[f"backbone.body.{suffix}"] = v
                    else:
                        remapped[k] = v
                checkpoint_state_dict = remapped
                print("  Detected MMRotate-style backbone keys; remapped backbone.* -> backbone.body.* for load.")

            ckpt_has_module_prefix = any(k.startswith("module.") for k in checkpoint_state_dict.keys())
            model_has_module_prefix = any(k.startswith("module.") for k in model_state_keys)

            # Handle common DDP/DataParallel prefix mismatch automatically.
            if ckpt_has_module_prefix and not model_has_module_prefix:
                checkpoint_state_dict = {
                    (k[7:] if k.startswith("module.") else k): v
                    for k, v in checkpoint_state_dict.items()
                }
                print("  Detected DDP checkpoint keys (module.*); stripping prefix for load.")
            elif model_has_module_prefix and not ckpt_has_module_prefix:
                checkpoint_state_dict = {f"module.{k}": v for k, v in checkpoint_state_dict.items()}
                print("  Model expects module.* keys; adding prefix to checkpoint state dict.")

            # Optional selective loading by prefix (config-driven transfer learning).
            include_prefixes = [p for p in (include_prefixes or []) if p]
            exclude_prefixes = [p for p in (exclude_prefixes or []) if p]
            if include_prefixes or exclude_prefixes:
                original_count = len(checkpoint_state_dict)
                filtered_state_dict: Dict[str, Any] = {}
                for key, value in checkpoint_state_dict.items():
                    key_no_module = key[7:] if key.startswith("module.") else key
                    if include_prefixes and not any(
                        key_no_module.startswith(prefix) for prefix in include_prefixes
                    ):
                        continue
                    if exclude_prefixes and any(
                        key_no_module.startswith(prefix) for prefix in exclude_prefixes
                    ):
                        continue
                    filtered_state_dict[key] = value
                checkpoint_state_dict = filtered_state_dict
                print(
                    "  Selective checkpoint loading enabled: "
                    f"kept {len(checkpoint_state_dict)}/{original_count} tensors "
                    f"(include={include_prefixes or None}, exclude={exclude_prefixes or None})."
                )

            # Drop keys whose tensor shapes differ from the model (e.g. ROI classifier when
            # num_classes changes, or RPN convs when anchor counts differ). PyTorch raises on
            # shape mismatch even when strict=False, which would otherwise abort the whole load.
            model_sd = model.state_dict()
            shape_skipped_detail: List[tuple] = []
            aligned_state_dict: Dict[str, Any] = {}
            for key, value in checkpoint_state_dict.items():
                if key not in model_sd:
                    continue
                if value.shape != model_sd[key].shape:
                    shape_skipped_detail.append(
                        (key, tuple(value.shape), tuple(model_sd[key].shape))
                    )
                    continue
                aligned_state_dict[key] = value
            if shape_skipped_detail:
                print(
                    f"  Skipped {len(shape_skipped_detail)} checkpoint tensors with shape mismatch "
                    f"(e.g. different num_classes or RPN anchor layout); model keeps init for those."
                )
                for key, ckpt_sh, model_sh in shape_skipped_detail[:10]:
                    print(f"     - {key}: checkpoint {ckpt_sh} vs model {model_sh}")
                if len(shape_skipped_detail) > 10:
                    print(f"     ... and {len(shape_skipped_detail) - 10} more")
            checkpoint_state_dict = aligned_state_dict

            missing_keys, unexpected_keys = model.load_state_dict(
                checkpoint_state_dict, strict=False
            )

            # Print warnings if there are mismatches
            if missing_keys:
                print(f"\n⚠️  Warning: Model architecture mismatch - {len(missing_keys)} keys not found in checkpoint:")
                print(f"   Missing keys (model expects but checkpoint doesn't have):")
                for key in missing_keys[:10]:  # Show first 10
                    print(f"     - {key}")
                if len(missing_keys) > 10:
                    print(f"     ... and {len(missing_keys) - 10} more")
                print(f"   These layers will be initialized randomly.")

            if unexpected_keys:
                print(f"\n⚠️  Warning: Model architecture mismatch - {len(unexpected_keys)} keys in checkpoint not used:")
                print(f"   Unexpected keys (checkpoint has but model doesn't need):")
                for key in unexpected_keys[:10]:  # Show first 10
                    print(f"     - {key}")
                if len(unexpected_keys) > 10:
                    print(f"     ... and {len(unexpected_keys) - 10} more")
                print(f"   These weights will be ignored.")

            if missing_keys or unexpected_keys:
                print(f"\n   Note: Training will continue with matching weights loaded and others initialized randomly.")
                print(f"   This is normal when changing model architecture (e.g., different num_classes, backbone, etc.)\n")

            # If strict=True was requested, raise error if there are mismatches
            if strict and (missing_keys or unexpected_keys):
                raise RuntimeError(
                    f"Strict loading failed: {len(missing_keys)} missing keys, "
                    f"{len(unexpected_keys)} unexpected keys. Set strict=False to allow partial loading."
                )

        except Exception as e:
            if strict:
                raise RuntimeError(f"Failed to load model state: {e}") from e
            else:
                print(f"\n⚠️  Warning: Failed to load model state: {e}")
                print(
                    "   Training will continue with default initialization for layers that "
                    "failed to load from the checkpoint.\n"
                )

        # Load optimizer state
        if optimizer is not None and "optimizer_state_dict" in checkpoint:
            try:
                optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
            except Exception as e:
                print(f"\n⚠️  Warning: Failed to load optimizer state: {e}")
                print(f"   Optimizer will be initialized with default state.\n")
        if scheduler is not None and "scheduler_state_dict" in checkpoint:
            try:
                scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
            except Exception as e:
                print(f"\n⚠️  Warning: Failed to load scheduler state: {e}")
                print(f"   Scheduler will be initialized with default state.\n")

        return checkpoint

    def get_latest(self) -> Optional[Path]:
        """Get path to latest checkpoint."""
        if not self.checkpoint_history:
            return None
        return self.checkpoint_history[-1]

    def get_best(self) -> Optional[Path]:
        """Get path to best checkpoint."""
        return self.best_checkpoint_path

get_best()

Get path to best checkpoint.

Source code in oriented_det/train/utils.py

def get_best(self) -> Optional[Path]:
    """Get path to best checkpoint."""
    return self.best_checkpoint_path

get_latest()

Get path to latest checkpoint.

Source code in oriented_det/train/utils.py

def get_latest(self) -> Optional[Path]:
    """Get path to latest checkpoint."""
    if not self.checkpoint_history:
        return None
    return self.checkpoint_history[-1]

load(checkpoint_path, model, optimizer=None, scheduler=None, *, strict=False, include_prefixes=None, exclude_prefixes=None)

Load a checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Path to checkpoint file	required
model	Module	Model to load weights into	required
optimizer	Optional[Optimizer]	Optional optimizer to load state into	None
strict	bool	If True, raise error on mismatch. If False, warn and load matching weights.	False

Returns:

Type	Description
Dict[str, Any]	Loaded checkpoint dict.

Source code in oriented_det/train/utils.py

def load(
    self,
    checkpoint_path: str | Path,
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    scheduler: Optional[Any] = None,
    *,
    strict: bool = False,
    include_prefixes: Optional[List[str]] = None,
    exclude_prefixes: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """Load a checkpoint.

    Args:
        checkpoint_path: Path to checkpoint file
        model: Model to load weights into
        optimizer: Optional optimizer to load state into
        strict: If True, raise error on mismatch. If False, warn and load matching weights.

    Returns:
        Loaded checkpoint dict.
    """
    _require_torch()

    checkpoint_path = Path(checkpoint_path)
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

    try:
        checkpoint = torch.load(checkpoint_path, map_location="cpu")
    except Exception as e:
        raise RuntimeError(f"Failed to load checkpoint: {e}") from e

    # Load model state dict
    try:
        # Support multiple checkpoint formats:
        # - this repo: {"model_state_dict": ...}
        # - MMDetection/MMRotate: {"state_dict": ...}
        # - raw state_dict: {param_name: tensor, ...}
        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
            checkpoint_state_dict = checkpoint["model_state_dict"]
        elif isinstance(checkpoint, dict) and "state_dict" in checkpoint:
            checkpoint_state_dict = checkpoint["state_dict"]
        elif isinstance(checkpoint, dict) and checkpoint and all(isinstance(k, str) for k in checkpoint.keys()):
            checkpoint_state_dict = checkpoint
        else:
            raise KeyError("Unsupported checkpoint format: expected model_state_dict or state_dict.")

        # MMRotate → torchvision BackboneWithFPN naming bridge:
        # MMRotate backbone keys are typically "backbone.*" (ResNet), while this repo's
        # torchvision backbone lives under "backbone.body.*".
        model_state_keys = list(model.state_dict().keys())
        model_has_backbone_body = any(k.startswith("backbone.body.") for k in model_state_keys)
        ckpt_has_backbone_body = any(k.startswith("backbone.body.") for k in checkpoint_state_dict.keys())
        ckpt_has_backbone = any(k.startswith("backbone.") for k in checkpoint_state_dict.keys())
        if model_has_backbone_body and ckpt_has_backbone and not ckpt_has_backbone_body:
            remapped = {}
            for k, v in checkpoint_state_dict.items():
                if k.startswith("backbone."):
                    suffix = k[len("backbone."):]
                    remapped[f"backbone.body.{suffix}"] = v
                else:
                    remapped[k] = v
            checkpoint_state_dict = remapped
            print("  Detected MMRotate-style backbone keys; remapped backbone.* -> backbone.body.* for load.")

        ckpt_has_module_prefix = any(k.startswith("module.") for k in checkpoint_state_dict.keys())
        model_has_module_prefix = any(k.startswith("module.") for k in model_state_keys)

        # Handle common DDP/DataParallel prefix mismatch automatically.
        if ckpt_has_module_prefix and not model_has_module_prefix:
            checkpoint_state_dict = {
                (k[7:] if k.startswith("module.") else k): v
                for k, v in checkpoint_state_dict.items()
            }
            print("  Detected DDP checkpoint keys (module.*); stripping prefix for load.")
        elif model_has_module_prefix and not ckpt_has_module_prefix:
            checkpoint_state_dict = {f"module.{k}": v for k, v in checkpoint_state_dict.items()}
            print("  Model expects module.* keys; adding prefix to checkpoint state dict.")

        # Optional selective loading by prefix (config-driven transfer learning).
        include_prefixes = [p for p in (include_prefixes or []) if p]
        exclude_prefixes = [p for p in (exclude_prefixes or []) if p]
        if include_prefixes or exclude_prefixes:
            original_count = len(checkpoint_state_dict)
            filtered_state_dict: Dict[str, Any] = {}
            for key, value in checkpoint_state_dict.items():
                key_no_module = key[7:] if key.startswith("module.") else key
                if include_prefixes and not any(
                    key_no_module.startswith(prefix) for prefix in include_prefixes
                ):
                    continue
                if exclude_prefixes and any(
                    key_no_module.startswith(prefix) for prefix in exclude_prefixes
                ):
                    continue
                filtered_state_dict[key] = value
            checkpoint_state_dict = filtered_state_dict
            print(
                "  Selective checkpoint loading enabled: "
                f"kept {len(checkpoint_state_dict)}/{original_count} tensors "
                f"(include={include_prefixes or None}, exclude={exclude_prefixes or None})."
            )

        # Drop keys whose tensor shapes differ from the model (e.g. ROI classifier when
        # num_classes changes, or RPN convs when anchor counts differ). PyTorch raises on
        # shape mismatch even when strict=False, which would otherwise abort the whole load.
        model_sd = model.state_dict()
        shape_skipped_detail: List[tuple] = []
        aligned_state_dict: Dict[str, Any] = {}
        for key, value in checkpoint_state_dict.items():
            if key not in model_sd:
                continue
            if value.shape != model_sd[key].shape:
                shape_skipped_detail.append(
                    (key, tuple(value.shape), tuple(model_sd[key].shape))
                )
                continue
            aligned_state_dict[key] = value
        if shape_skipped_detail:
            print(
                f"  Skipped {len(shape_skipped_detail)} checkpoint tensors with shape mismatch "
                f"(e.g. different num_classes or RPN anchor layout); model keeps init for those."
            )
            for key, ckpt_sh, model_sh in shape_skipped_detail[:10]:
                print(f"     - {key}: checkpoint {ckpt_sh} vs model {model_sh}")
            if len(shape_skipped_detail) > 10:
                print(f"     ... and {len(shape_skipped_detail) - 10} more")
        checkpoint_state_dict = aligned_state_dict

        missing_keys, unexpected_keys = model.load_state_dict(
            checkpoint_state_dict, strict=False
        )

        # Print warnings if there are mismatches
        if missing_keys:
            print(f"\n⚠️  Warning: Model architecture mismatch - {len(missing_keys)} keys not found in checkpoint:")
            print(f"   Missing keys (model expects but checkpoint doesn't have):")
            for key in missing_keys[:10]:  # Show first 10
                print(f"     - {key}")
            if len(missing_keys) > 10:
                print(f"     ... and {len(missing_keys) - 10} more")
            print(f"   These layers will be initialized randomly.")

        if unexpected_keys:
            print(f"\n⚠️  Warning: Model architecture mismatch - {len(unexpected_keys)} keys in checkpoint not used:")
            print(f"   Unexpected keys (checkpoint has but model doesn't need):")
            for key in unexpected_keys[:10]:  # Show first 10
                print(f"     - {key}")
            if len(unexpected_keys) > 10:
                print(f"     ... and {len(unexpected_keys) - 10} more")
            print(f"   These weights will be ignored.")

        if missing_keys or unexpected_keys:
            print(f"\n   Note: Training will continue with matching weights loaded and others initialized randomly.")
            print(f"   This is normal when changing model architecture (e.g., different num_classes, backbone, etc.)\n")

        # If strict=True was requested, raise error if there are mismatches
        if strict and (missing_keys or unexpected_keys):
            raise RuntimeError(
                f"Strict loading failed: {len(missing_keys)} missing keys, "
                f"{len(unexpected_keys)} unexpected keys. Set strict=False to allow partial loading."
            )

    except Exception as e:
        if strict:
            raise RuntimeError(f"Failed to load model state: {e}") from e
        else:
            print(f"\n⚠️  Warning: Failed to load model state: {e}")
            print(
                "   Training will continue with default initialization for layers that "
                "failed to load from the checkpoint.\n"
            )

    # Load optimizer state
    if optimizer is not None and "optimizer_state_dict" in checkpoint:
        try:
            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        except Exception as e:
            print(f"\n⚠️  Warning: Failed to load optimizer state: {e}")
            print(f"   Optimizer will be initialized with default state.\n")
    if scheduler is not None and "scheduler_state_dict" in checkpoint:
        try:
            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        except Exception as e:
            print(f"\n⚠️  Warning: Failed to load scheduler state: {e}")
            print(f"   Scheduler will be initialized with default state.\n")

    return checkpoint

save(model, optimizer=None, scheduler=None, epoch=None, metrics=None, *, suffix='', custom_filename=None)

Save a checkpoint.

If custom_filename is set (e.g. "best_mAP_0.42.pth"), that name is used and the checkpoint is not added to history (treated as a best checkpoint).

Returns path to saved checkpoint.

Source code in oriented_det/train/utils.py

def save(
    self,
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    scheduler: Optional[Any] = None,
    epoch: Optional[int] = None,
    metrics: Optional[Dict[str, float]] = None,
    *,
    suffix: str = "",
    custom_filename: Optional[str] = None,
) -> Path:
    """Save a checkpoint.

    If custom_filename is set (e.g. "best_mAP_0.42.pth"), that name is used
    and the checkpoint is not added to history (treated as a best checkpoint).

    Returns path to saved checkpoint.
    """
    _require_torch()
    self.checkpoint_dir.mkdir(parents=True, exist_ok=True)

    checkpoint = {
        "model_state_dict": model.state_dict(),
        "epoch": epoch,
        "metrics": metrics or {},
    }

    if optimizer is not None:
        checkpoint["optimizer_state_dict"] = optimizer.state_dict()
    if scheduler is not None and hasattr(scheduler, "state_dict"):
        checkpoint["scheduler_state_dict"] = scheduler.state_dict()

    if custom_filename:
        filename = custom_filename if custom_filename.endswith(".pth") else f"{custom_filename}.pth"
    elif suffix:
        filename = f"checkpoint_{suffix}.pth"
    elif epoch is not None:
        filename = f"checkpoint_epoch_{epoch:04d}.pth"
    else:
        filename = f"checkpoint_{int(time.time())}.pth"

    checkpoint_path = self.checkpoint_dir / filename

    # Atomic write: write to temp file first, then rename. Retry on transient I/O errors
    # (e.g. NFS "unexpected pos" / "file write failed"). Fallback: write to local /tmp then move.
    temp_path = checkpoint_path.with_suffix(".tmp")
    max_attempts = 3
    last_error = None
    for attempt in range(max_attempts):
        try:
            torch.save(checkpoint, temp_path)
            temp_path.replace(checkpoint_path)
            last_error = None
            break
        except Exception as e:
            last_error = e
            if temp_path.exists():
                try:
                    temp_path.unlink()
                except OSError:
                    pass
            if attempt < max_attempts - 1:
                time.sleep(1.0 * (attempt + 1))
                continue
    if last_error is not None:
        # Fallback: write to local temp dir (avoids NFS/network fs write failures)
        local_tmp = Path(os.environ.get("TMPDIR", "/tmp"))
        if local_tmp.exists() and os.access(local_tmp, os.W_OK):
            fallback_temp = local_tmp / f"oriented_det_checkpoint_{os.getpid()}_{filename}.tmp"
            try:
                torch.save(checkpoint, fallback_temp)
                shutil.move(str(fallback_temp), str(checkpoint_path))
            except Exception as e2:
                if fallback_temp.exists():
                    try:
                        fallback_temp.unlink()
                    except OSError:
                        pass
                raise RuntimeError(
                    f"Failed to save checkpoint after {max_attempts} attempts and local fallback: {last_error!s}; fallback: {e2!s}"
                ) from e2
        else:
            raise RuntimeError(f"Failed to save checkpoint: {last_error}") from last_error

    # Only add to history if it's not a best checkpoint (best is managed separately)
    if not custom_filename and suffix != "best":
        self.checkpoint_history.append(checkpoint_path)

        # Clean up old checkpoints (excluding best checkpoint)
        while len(self.checkpoint_history) > self.keep_last_n:
            old_checkpoint = self.checkpoint_history.pop(0)
            # Don't delete if it's the best checkpoint
            if old_checkpoint != self.best_checkpoint_path and old_checkpoint.exists():
                old_checkpoint.unlink()

    return checkpoint_path

save_best(model, optimizer=None, scheduler=None, epoch=None, metrics=None)

Save checkpoint if it's the best according to best_metric.

The best checkpoint is stored independently and is not subject to the keep_last_n cleanup, ensuring it's always preserved.

Source code in oriented_det/train/utils.py

def save_best(
    self,
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    scheduler: Optional[Any] = None,
    epoch: Optional[int] = None,
    metrics: Optional[Dict[str, float]] = None,
) -> Optional[Path]:
    """Save checkpoint if it's the best according to best_metric.

    The best checkpoint is stored independently and is not subject to
    the keep_last_n cleanup, ensuring it's always preserved.
    """
    if self.best_metric is None or metrics is None:
        return None

    current_value = metrics.get(self.best_metric)
    if current_value is None:
        return None
    # Do not save best checkpoint when mAP was not computed (sentinel -1)
    if str(self.best_metric).strip().lower() == "map" and current_value == -1.0:
        return None

    is_better = (
        (self.best_value is None) or
        (self.higher_is_better and current_value > self.best_value) or
        (not self.higher_is_better and current_value < self.best_value)
    )

    if is_better:
        # Delete old best checkpoint if it exists
        if self.best_checkpoint_path is not None and self.best_checkpoint_path.exists():
            try:
                self.best_checkpoint_path.unlink()
            except Exception as e:
                # Log but don't fail if deletion fails
                print(f"Warning: Failed to delete old best checkpoint {self.best_checkpoint_path}: {e}")

        self.best_value = current_value
        # Build filename with metric name and value, e.g. best_mAP_0.42.pth or best_total_loss_1.23.pth
        safe_metric = "".join(c if c.isalnum() or c == "_" else "_" for c in str(self.best_metric))
        if isinstance(current_value, float):
            value_str = f"{current_value:.2f}" if safe_metric.lower() == "map" else f"{current_value:.4f}"
        else:
            value_str = str(current_value)
        custom_filename = f"best_{safe_metric}_{value_str}.pth"
        self.best_checkpoint_path = self.save(
            model, optimizer, scheduler, epoch, metrics, custom_filename=custom_filename
        )
        return self.best_checkpoint_path

    return None

CosineAnnealingWithFixedTailLR

Bases: _lr_scheduler_base_class()

Cosine decay for cosine_epochs scheduler steps, then a constant LR tail.

Matches :class:torch.optim.lr_scheduler.CosineAnnealingLR for last_epoch in [1, cosine_epochs] (same closed form as PyTorch's _get_closed_form_lr), then holds tail_lr for last_epoch > cosine_epochs. The last_epoch == 0 branch keeps the current optimizer LR (same as PyTorch cosine) so it composes with epoch-based stepping after warmup.

cosine_epochs corresponds to T_max in CosineAnnealingLR (minimum LR is reached when last_epoch == cosine_epochs).

Source code in oriented_det/train/utils.py

class CosineAnnealingWithFixedTailLR(_lr_scheduler_base_class()):  # type: ignore[misc, valid-type]
    """Cosine decay for ``cosine_epochs`` scheduler steps, then a constant LR tail.

    Matches :class:`torch.optim.lr_scheduler.CosineAnnealingLR` for
    ``last_epoch in [1, cosine_epochs]`` (same closed form as PyTorch's ``_get_closed_form_lr``),
    then holds ``tail_lr`` for ``last_epoch > cosine_epochs``. The ``last_epoch == 0`` branch keeps
    the current optimizer LR (same as PyTorch cosine) so it composes with epoch-based stepping after warmup.

    ``cosine_epochs`` corresponds to ``T_max`` in ``CosineAnnealingLR`` (minimum LR is reached when
    ``last_epoch == cosine_epochs``).
    """

    def __init__(
        self,
        optimizer: "optim.Optimizer",
        cosine_epochs: int,
        eta_min: float = 0.0,
        tail_lr: Optional[float] = None,
        last_epoch: int = -1,
    ) -> None:
        _require_torch()
        self.cosine_epochs = int(cosine_epochs)
        if self.cosine_epochs < 1:
            raise ValueError("cosine_epochs must be >= 1")
        self.eta_min = float(eta_min)
        self.tail_lr = float(self.eta_min if tail_lr is None else tail_lr)
        super().__init__(optimizer, last_epoch)

    def get_lr(self) -> List[float]:
        if self.last_epoch == 0:
            return [group["lr"] for group in self.optimizer.param_groups]
        if self.last_epoch <= self.cosine_epochs:
            return [
                self.eta_min
                + (base_lr - self.eta_min)
                * (1.0 + math.cos(math.pi * float(self.last_epoch) / float(self.cosine_epochs)))
                / 2.0
                for base_lr in self.base_lrs
            ]
        return [self.tail_lr for _ in self.base_lrs]

DataLoaderConfig dataclass

DataLoader configuration.

Source code in oriented_det/train/config.py

@dataclass
class DataLoaderConfig:
    """DataLoader configuration."""
    batch_size: int = 16
    num_workers: int = 4
    shuffle: bool = True
    pin_memory: bool = True

DatasetConfig dataclass

Dataset configuration.

Source code in oriented_det/train/config.py

@dataclass
class DatasetConfig:
    """Dataset configuration."""
    data_root: Path
    format: str = "dota"  # Options: "dota", "airbus_playground"
    train_tiles_dir: Optional[Path] = None
    val_tiles_dir: Optional[Path] = None
    # Optional list of tile roots (MMRotate trainval-style union without on-disk merge).
    train_tiles_dirs: Optional[List[Path]] = None
    val_tiles_dirs: Optional[List[Path]] = None
    same_folder: bool = False  # If True (DOTA), images and labels are in the same directory as train_tiles_dir/val_tiles_dir
    overlap: int = 16  # Tile overlap in pixels (even); deploy uses margin = overlap/2 when production.ignore_margin_pixels is null
    annotations_file: Optional[Path] = None
    split_file: Optional[Path] = None
    # Airbus Playground split.csv: integer fold id used as validation (default 0). Ignored for
    # legacy CSVs whose split column uses only the strings train/val.
    val_split_id: int = 0
    # Airbus Playground fold mode: when true, train split includes all folds (DOTA trainval parity);
    # val split still uses val_split_id for monitoring only.
    train_includes_val: bool = False
    ignore_labels: Optional[List[str]] = None
    map_labels: Optional[Dict[str, str]] = None
    # How to handle DOTA "difficult" annotations (last field in label lines):
    # - "drop": remove difficult objects at read-time (fast; they never reach training/eval targets)
    # - "ignore": keep difficult objects but route them into ignore targets (MMRotate/MMDet-style)
    # - "keep": treat difficult objects as normal GT
    difficult_strategy: str = "drop"
    # Drop tiles with no GT after difficult_strategy / allowed_classes / ignore_labels (MMRotate parity).
    filter_empty_gt: bool = False
    max_train_samples: Optional[int] = None
    max_val_samples: Optional[int] = None
    # When set with max_train_samples / max_val_samples: pick indices via deterministic shuffle
    # instead of the first N in dataset order (see capped_subset_indices in train.utils).
    max_samples_shuffle_seed: Optional[int] = None
    allowed_classes: Optional[List[str]] = None
    # Optional CSV from tools/save_predictions (--save-tile-metrics-csv); join on image_id / stem
    tile_metrics_csv: Optional[Path] = None
    hard_tile_metric_column: str = "f1"
    hard_tile_threshold: float = 0.8
    hard_tile_oversample_factor: float = 2.0

    def __post_init__(self):
        """Convert string paths to Path objects."""
        self.data_root = Path(self.data_root)
        if self.train_tiles_dir is not None:
            self.train_tiles_dir = Path(self.train_tiles_dir)
        if self.val_tiles_dir is not None:
            self.val_tiles_dir = Path(self.val_tiles_dir)
        if self.train_tiles_dirs is not None:
            self.train_tiles_dirs = [Path(p) for p in self.train_tiles_dirs]
        if self.val_tiles_dirs is not None:
            self.val_tiles_dirs = [Path(p) for p in self.val_tiles_dirs]
        if self.annotations_file is not None:
            self.annotations_file = Path(self.annotations_file)
        if self.split_file is not None:
            self.split_file = Path(self.split_file)
        if self.tile_metrics_csv is not None:
            self.tile_metrics_csv = Path(self.tile_metrics_csv)

    def has_dota_tiles_config(self) -> bool:
        """True when train and val tile root(s) are configured."""
        train_ok = bool(self.train_tiles_dirs) or self.train_tiles_dir is not None
        val_ok = bool(self.val_tiles_dirs) or self.val_tiles_dir is not None
        return train_ok and val_ok

    def get_train_tile_roots(self) -> List[Path]:
        from ..data.dota import resolve_dota_tile_roots

        return resolve_dota_tile_roots(
            tiles_dirs=self.train_tiles_dirs,
            tiles_dir=self.train_tiles_dir,
            split_label="train",
        )

    def get_val_tile_roots(self) -> List[Path]:
        from ..data.dota import resolve_dota_tile_roots

        return resolve_dota_tile_roots(
            tiles_dirs=self.val_tiles_dirs,
            tiles_dir=self.val_tiles_dir,
            split_label="val",
        )

__post_init__()

Convert string paths to Path objects.

Source code in oriented_det/train/config.py

def __post_init__(self):
    """Convert string paths to Path objects."""
    self.data_root = Path(self.data_root)
    if self.train_tiles_dir is not None:
        self.train_tiles_dir = Path(self.train_tiles_dir)
    if self.val_tiles_dir is not None:
        self.val_tiles_dir = Path(self.val_tiles_dir)
    if self.train_tiles_dirs is not None:
        self.train_tiles_dirs = [Path(p) for p in self.train_tiles_dirs]
    if self.val_tiles_dirs is not None:
        self.val_tiles_dirs = [Path(p) for p in self.val_tiles_dirs]
    if self.annotations_file is not None:
        self.annotations_file = Path(self.annotations_file)
    if self.split_file is not None:
        self.split_file = Path(self.split_file)
    if self.tile_metrics_csv is not None:
        self.tile_metrics_csv = Path(self.tile_metrics_csv)

has_dota_tiles_config()

True when train and val tile root(s) are configured.

Source code in oriented_det/train/config.py

def has_dota_tiles_config(self) -> bool:
    """True when train and val tile root(s) are configured."""
    train_ok = bool(self.train_tiles_dirs) or self.train_tiles_dir is not None
    val_ok = bool(self.val_tiles_dirs) or self.val_tiles_dir is not None
    return train_ok and val_ok

EvaluationConfig dataclass

Evaluation configuration.

Source code in oriented_det/train/config.py

@dataclass
class EvaluationConfig:
    """Evaluation configuration."""
    score_threshold: float = 0.5  # Minimum confidence score for detections
    # Optional class_name -> min score; classes not listed use score_threshold
    per_class_score_threshold: Optional[Dict[str, float]] = None
    iou_threshold: float = 0.5  # IoU threshold for mAP calculation
    compute_map_final: bool = True  # After training, load best checkpoint and compute mAP once
    compute_map_every_n_epochs: int = 0  # If >0, compute mAP every N epochs during training (current model)
    # If True, compute expensive GT–IoU histograms (mean/median/buckets, wrong-class counts)
    # and class-agnostic duplicate rate on post-threshold detections (greedy by score vs eval IoU).
    extended_gt_metrics: bool = False
    # mAP matching and GT-cover / accuracy use exact CPU polygon IoU when True (Shapely when
    # installed). When False, use GPU sampling IoU (faster, approximate). Does not affect final
    # detection NMS or production/deploy decode settings.
    use_exact_rotated_iou: bool = True
    # IoU backend for compute_map_final only (best-checkpoint mAP after training). When None,
    # falls back to use_exact_rotated_iou. Set true for publishable exact mAP while keeping
    # periodic in-training mAP on GPU sampling (use_exact_rotated_iou=false).
    use_exact_rotated_iou_for_final_map: Optional[bool] = None

GroupedCeSpec dataclass

Resolved group membership for ROI classifier curriculum.

Source code in oriented_det/train/grouped_ce.py

@dataclass(frozen=True)
class GroupedCeSpec:
    """Resolved group membership for ROI classifier curriculum."""

    group_index_lists: Tuple[Tuple[int, ...], ...]  # 1-indexed foreground class ids per group
    class_in_group_id: Tuple[int, ...]  # length num_classes+1; -1 bg/unmapped, else group index

LossConfig dataclass

Experiment-level ROI loss selection and optional class weighting.

loss_type chooses the recipe applied by tools/train.py: - cross_entropy: plain ROI CE (MMRotate-style) - class_weighted: ROI CE with dataset-derived class weights - focal: unweighted focal ROI loss - focal_weighted: focal ROI loss with dataset-derived class weights - none: legacy fallback to model.roi_loss_type

Source code in oriented_det/train/config.py

@dataclass
class LossConfig:
    """Experiment-level ROI loss selection and optional class weighting.

    loss_type chooses the recipe applied by tools/train.py:
    - cross_entropy: plain ROI CE (MMRotate-style)
    - class_weighted: ROI CE with dataset-derived class weights
    - focal: unweighted focal ROI loss
    - focal_weighted: focal ROI loss with dataset-derived class weights
    - none: legacy fallback to model.roi_loss_type
    """
    loss_type: str = "class_weighted"
    class_weight_method: str = "sqrt"
    class_weight_beta: float = 0.9999  # Used by effective_num weighting
    # Optional schedule to ramp class weights in over epochs (e.g. uniform -> inv_freq).
    class_weight_schedule_type: Optional[str] = None  # None | "linear_ramp"
    class_weight_schedule_start_epoch: int = 0
    class_weight_schedule_end_epoch: int = 0
    class_weight_schedule_power: float = 1.0  # 1.0 = linear; >1 eases in
    background_weight: Optional[float] = None
    focal_alpha: float = 1.0
    focal_gamma: float = 2.0
    class_weight_overrides: Optional[Dict[str, float]] = None  # e.g. {"truck": 0.25} to reduce dominant-class bias
    label_smoothing: float = 0.0  # Label smoothing for ROI classification (e.g. 0.1 to reduce overconfidence)
    # Coarse-to-fine ROI classifier curriculum (single-run alternative to multi-stage training).
    roi_grouped_ce_enabled: bool = False
    roi_grouped_ce_groups: Optional[Dict[str, List[str]]] = None  # e.g. {"plane": ["Bomber Aircraft", ...]}
    roi_grouped_ce_schedule_type: Optional[str] = None  # None | "step" | "linear_ramp"
    roi_grouped_ce_schedule_start_epoch: int = 0
    roi_grouped_ce_schedule_end_epoch: int = 8  # step: first fine epoch; ramp: last grouped-heavy epoch
    roi_grouped_ce_schedule_power: float = 1.0

MetricTracker dataclass

Efficient metric tracking during training.

Source code in oriented_det/train/utils.py

@dataclass
class MetricTracker:
    """Efficient metric tracking during training."""

    def __init__(self):
        self.metrics: Dict[str, List[float]] = defaultdict(list)
        self.step_times: List[float] = []

    def update(self, metrics: Dict[str, float]) -> None:
        """Update metrics with new values."""
        for key, value in metrics.items():
            self.metrics[key].append(float(value))

    def update_time(self, elapsed: float) -> None:
        """Record step time."""
        self.step_times.append(elapsed)

    def get_average(self, key: str, window: Optional[int] = None) -> float:
        """Get average of a metric, optionally over a window."""
        values = self.metrics[key]
        if not values:
            return 0.0
        if window is None:
            return sum(values) / len(values)
        recent = values[-window:]
        return sum(recent) / len(recent)

    def get_latest(self, key: str) -> Optional[float]:
        """Get most recent value of a metric."""
        values = self.metrics[key]
        return values[-1] if values else None

    def get_summary(self, window: Optional[int] = None) -> Dict[str, float]:
        """Get summary of all metrics."""
        summary = {}
        for key in self.metrics:
            summary[key] = self.get_average(key, window=window)
        if self.step_times:
            recent_times = self.step_times[-window:] if window else self.step_times
            summary["time_per_step"] = sum(recent_times) / len(recent_times)
        return summary

    def reset(self) -> None:
        """Reset all metrics."""
        self.metrics.clear()
        self.step_times.clear()

get_average(key, window=None)

Get average of a metric, optionally over a window.

Source code in oriented_det/train/utils.py

def get_average(self, key: str, window: Optional[int] = None) -> float:
    """Get average of a metric, optionally over a window."""
    values = self.metrics[key]
    if not values:
        return 0.0
    if window is None:
        return sum(values) / len(values)
    recent = values[-window:]
    return sum(recent) / len(recent)

get_latest(key)

Get most recent value of a metric.

Source code in oriented_det/train/utils.py

def get_latest(self, key: str) -> Optional[float]:
    """Get most recent value of a metric."""
    values = self.metrics[key]
    return values[-1] if values else None

get_summary(window=None)

Get summary of all metrics.

Source code in oriented_det/train/utils.py

def get_summary(self, window: Optional[int] = None) -> Dict[str, float]:
    """Get summary of all metrics."""
    summary = {}
    for key in self.metrics:
        summary[key] = self.get_average(key, window=window)
    if self.step_times:
        recent_times = self.step_times[-window:] if window else self.step_times
        summary["time_per_step"] = sum(recent_times) / len(recent_times)
    return summary

reset()

Reset all metrics.

Source code in oriented_det/train/utils.py

def reset(self) -> None:
    """Reset all metrics."""
    self.metrics.clear()
    self.step_times.clear()

update(metrics)

Update metrics with new values.

Source code in oriented_det/train/utils.py

def update(self, metrics: Dict[str, float]) -> None:
    """Update metrics with new values."""
    for key, value in metrics.items():
        self.metrics[key].append(float(value))

update_time(elapsed)

Record step time.

Source code in oriented_det/train/utils.py

def update_time(self, elapsed: float) -> None:
    """Record step time."""
    self.step_times.append(elapsed)

ModelConfig dataclass

Model configuration. Keys ordered by prefix: backbone, fpn_, anchor_, target_, roi_, rpn_, final_nms_*, etc.

Source code in oriented_det/train/config.py

@dataclass
class ModelConfig:
    """Model configuration. Keys ordered by prefix: backbone, fpn_, anchor_, target_, roi_*, rpn_*, final_nms_*, etc."""
    # Backbone
    backbone: str = "resnet50"
    pretrained_backbone: bool = True
    trainable_layers: int = 5
    frozen_stages: Optional[int] = None
    # FPN
    fpn_returned_layers: Optional[List[int]] = None
    # Nominal strides (ROI init / logging); RPN/ROI forward uses strides derived from image_size and feature maps
    fpn_strides: Optional[List[int]] = None
    fpn_extra_level: bool = False
    # Anchors
    anchor_scales: List[int] = field(default_factory=lambda: [8, 16, 32])
    anchor_ratios: List[float] = field(default_factory=lambda: [0.5, 1.0, 2.0])
    anchor_octave_base_scale: Optional[float] = None
    anchor_scales_per_octave: Optional[int] = None
    # Box regression targets [dx, dy, dw, dh, dangle]
    target_means: tuple[float, float, float, float, float] = (0.0, 0.0, 0.0, 0.0, 0.0)
    target_stds: tuple[float, float, float, float, float] = (0.1, 0.1, 0.2, 0.2, 0.1)
    # ROI head (switch + params)
    roi_loss_type: str = "cross_entropy"
    roi_focal_alpha: float = 1.0
    roi_focal_gamma: float = 2.0
    roi_norm_factor: Optional[float] = 2.0
    roi_edge_swap: bool = True
    roi_proj_xy: bool = False
    roi_box_reg_angle_weight: float = 1.0
    roi_box_reg_angle_schedule_epochs: Optional[List[int]] = None
    roi_box_reg_angle_schedule_values: Optional[List[float]] = None
    roi_box_reg_iou_weight: float = 0.0
    roi_box_reg_iou_schedule_epochs: Optional[List[int]] = None
    roi_box_reg_iou_schedule_values: Optional[List[float]] = None
    roi_batch_size_per_image: int = 512
    roi_positive_iou_threshold: float = 0.5
    roi_negative_iou_threshold: float = 0.5
    roi_match_low_quality: bool = False
    roi_min_pos_iou: float = 0.5
    roi_box_reg_iou_loss_type: str = "riou"
    roi_box_reg_kfiou_fun: Optional[str] = None
    roi_box_reg_probiou_mode: Optional[str] = None
    roi_box_reg_main_loss_type: str = "smooth_l1"
    roi_box_reg_norm: str = "sampled_all"
    roi_box_reg_smooth_l1_aux_weight: float = 0.0
    # RetinaNet head / regression
    retinanet_stacked_convs: int = 4
    box_reg_loss_type: str = "smooth_l1"
    box_reg_weight: float = 1.0
    # RPN
    rpn_min_size: float = 0.0
    rpn_pre_nms_top_n: int = 2000
    rpn_post_nms_top_n: int = 1000
    rpn_nms_threshold: float = 0.7
    rpn_batch_size_per_image: int = 256
    rpn_positive_iou_threshold: float = 0.7
    rpn_negative_iou_threshold: float = 0.3
    rpn_min_pos_iou: float = 0.3
    rpn_match_low_quality: bool = True
    # Matching / proposals
    use_hbb_for_matching: bool = True
    roi_use_hbb_for_matching: bool = False
    add_gt_as_proposals: bool = True
    # NMS and inference (after ROI head; distinct from rpn_nms_threshold)
    final_nms_iou_threshold: float = 0.5
    # If True, final ROI NMS runs on all classes together (reduces overlapping car/truck duplicates).
    nms_class_agnostic: bool = False
    # If True, post-ROI final oriented NMS uses exact polygon IoU on CPU (ignores GPU sampling NMS).
    final_nms_use_cpu: bool = False
    final_nms_iou_schedule_epochs: Optional[List[int]] = None
    final_nms_iou_schedule_values: Optional[List[float]] = None
    max_detections_per_image: int = 100
    inference_pre_nms_score_threshold: float = 0.05
    # If True, ROI inference uses argmax foreground class per proposal (MMRotate-style); if False,
    # keep every foreground class above inference_pre_nms_score_threshold (helps weak classifiers).
    roi_inference_top_class_only: bool = False

OneCycleWrapper

Wrapper for OneCycleLR so it is stepped every optimizer step (like WarmupScheduler).

OneCycleLR must be stepped each batch; the main train loop only steps schedulers once per epoch. This wrapper exposes step() (forwarded to OneCycleLR) and step_epoch() (no-op), so train_one_epoch can call step() every batch and the main loop's step_epoch() does nothing.

Source code in oriented_det/train/utils.py

class OneCycleWrapper:
    """Wrapper for OneCycleLR so it is stepped every optimizer step (like WarmupScheduler).

    OneCycleLR must be stepped each batch; the main train loop only steps schedulers
    once per epoch. This wrapper exposes step() (forwarded to OneCycleLR) and
    step_epoch() (no-op), so train_one_epoch can call step() every batch and the
    main loop's step_epoch() does nothing.
    """

    def __init__(self, one_cycle_scheduler: Any) -> None:
        self._scheduler = one_cycle_scheduler

    def step(self) -> None:
        self._scheduler.step()

    def step_epoch(self) -> None:
        # No-op: OneCycleLR is stepped per batch only
        pass

    def state_dict(self) -> Dict[str, Any]:
        """Get state dict for checkpointing."""
        return self._scheduler.state_dict()

    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
        """Load state dict from checkpoint."""
        self._scheduler.load_state_dict(state_dict)

load_state_dict(state_dict)

Load state dict from checkpoint.

Source code in oriented_det/train/utils.py

def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
    """Load state dict from checkpoint."""
    self._scheduler.load_state_dict(state_dict)

state_dict()

Get state dict for checkpointing.

Source code in oriented_det/train/utils.py

def state_dict(self) -> Dict[str, Any]:
    """Get state dict for checkpointing."""
    return self._scheduler.state_dict()

ProductionConfig dataclass

Production overrides (deploy, save_predictions, image_demo).

None on a field means: use the same source as before this block existed (evaluation / model / dataset), or deploy defaults for canvas flags (see schema).

final_nms_iou_threshold / final_nms_use_cpu mirror model.* when set here; decode patches apply only when apply_inference_config_to_model runs on a loaded checkpoint (not during tools/train.py).

Source code in oriented_det/train/config.py

@dataclass
class ProductionConfig:
    """Production overrides (deploy, save_predictions, image_demo).

    ``None`` on a field means: use the same source as before this block existed
    (``evaluation`` / ``model`` / ``dataset``), or deploy defaults for canvas flags (see schema).

    ``final_nms_iou_threshold`` / ``final_nms_use_cpu`` mirror ``model.*`` when set here;
    decode patches apply only when ``apply_inference_config_to_model`` runs on a loaded checkpoint
    (not during ``tools/train.py``).
    """

    score_threshold: Optional[float] = None
    per_class_score_threshold: Optional[Dict[str, float]] = None
    final_nms_iou_threshold: Optional[float] = None
    final_nms_use_cpu: Optional[bool] = None
    inference_pre_nms_score_threshold: Optional[float] = None
    max_detections_per_image: Optional[int] = None
    nms_class_agnostic: Optional[bool] = None
    roi_inference_top_class_only: Optional[bool] = None
    rpn_pre_nms_top_n: Optional[int] = None
    rpn_post_nms_top_n: Optional[int] = None
    rpn_nms_threshold: Optional[float] = None
    # Sliding-window overlap along each axis (pixels). None = 200 (same default as oriented_det.runtime.inference).
    overlap_pixels: Optional[int] = None
    # Interior margin (px): drop detections whose centroid lies in the outer band. None = dataset.overlap/2 (0 = off).
    ignore_margin_pixels: Optional[float] = None
    # Deploy: expand square canvas from first image (fixed resize). None = false.
    use_first_image_canvas: Optional[bool] = None
    # Deploy: always use config target_size / sliding window. None = true.
    stick_to_model_canvas: Optional[bool] = None

TrainingConfig dataclass

Training hyperparameters. Keys ordered: scheduler switch first, then lr_*, then rest.

Source code in oriented_det/train/config.py

@dataclass
class TrainingConfig:
    """Training hyperparameters. Keys ordered: scheduler switch first, then lr_*, then rest."""
    # LR scheduler: switch first, then all lr_* grouped
    lr_scheduler_type: Optional[str] = None  # None/"multistep"/"step" | "reduce_on_plateau" | "one_cycle" | "cosine_annealing" | "cosine_annealing_with_tail"
    lr_scheduler_step_epochs: int = 8
    lr_scheduler_milestones: Optional[List[int]] = None
    lr_scheduler_gamma: float = 0.1
    lr_scheduler_plateau_metric: str = "total_loss"
    lr_scheduler_plateau_factor: float = 0.1
    lr_scheduler_plateau_patience: int = 5
    lr_scheduler_one_cycle_pct_start: float = 0.3
    lr_scheduler_one_cycle_div_factor: float = 25.0
    lr_scheduler_one_cycle_final_div_factor: float = 1e4
    lr_scheduler_cosine_t_max: Optional[int] = None
    lr_scheduler_cosine_eta_min: float = 1e-6
    # Cosine phase length (epochs). If set with num_epochs > this value, remaining epochs use a
    # fixed LR tail (see lr_scheduler_cosine_tail_lr). Alternative: set lr_scheduler_cosine_tail_epochs only.
    lr_scheduler_cosine_epochs: Optional[int] = None
    lr_scheduler_cosine_tail_epochs: int = 0
    lr_scheduler_cosine_tail_lr: Optional[float] = None  # default: lr_scheduler_cosine_eta_min
    lr_warmup_steps: int = 500
    lr_scaling_with_accumulation: str = "linear"
    lr_scale_with_world_size: bool = False
    use_lr_param_groups: bool = False
    lr_mult_backbone: float = 0.5
    lr_mult_rpn: float = 0.5
    lr_mult_roi: float = 2.0
    lr_mult_other: float = 1.0
    # When set: applies to RetinaNet ``head.*``; for two-stage models, the same value applies to
    # ``rpn_head.*`` and ``roi_head.*`` (overrides lr_mult_rpn / lr_mult_roi for those groups).
    # When null: ``head.*`` uses lr_mult_other; RPN/ROI use lr_mult_rpn / lr_mult_roi.
    lr_mult_head: Optional[float] = None
    # Other training
    num_epochs: int = 20
    learning_rate: float = 0.001
    momentum: float = 0.9
    weight_decay: float = 0.0001
    use_amp: bool = True
    gradient_accumulation_steps: int = 2
    max_grad_norm: float = 1.0
    loss_weights: Optional[Dict[str, float]] = None
    # Optional partial freeze (0-based training-loop epoch index ``epoch``; same as checkpoint epoch).
    # While ``epoch < N``, the corresponding module does not train (``requires_grad=False``). ROI head
    # always remains trainable. RPN freeze is ignored when the model has no ``rpn_head`` (e.g. RetinaNet).
    freeze_backbone_epochs: int = 0
    freeze_rpn_epochs: int = 0
    # Early stopping (optional). When early_stop_patience is set, stop if monitored metric does not improve.
    early_stop_patience: Optional[int] = None
    early_stop_metric: str = "mAP"
    early_stop_min_delta: float = 0.0
    early_stop_higher_is_better: Optional[bool] = None

TrainingExperimentConfig dataclass

Complete training experiment configuration.

Source code in oriented_det/train/config.py

@dataclass
class TrainingExperimentConfig:
    """Complete training experiment configuration."""
    # Experiment metadata
    model_type: str = "oriented_rcnn"
    experiment_timestamp: Optional[str] = None
    source_recipe: Optional[str] = None
    # Set at train start (not in hand-written recipes): framework git + package stamp
    source_code_root: Optional[str] = None
    git_commit: Optional[str] = None
    git_describe: Optional[str] = None
    git_dirty: Optional[bool] = None
    git_branch: Optional[str] = None
    git_commit_date: Optional[str] = None
    package_version: Optional[str] = None
    # Switches first
    enable_albumentation: bool = False
    enable_profiling: bool = False
    # Sub-configs (grouped by section)
    dataset: Optional[DatasetConfig] = None
    augmentation: AugmentationConfig = field(default_factory=AugmentationConfig)
    data_loader: DataLoaderConfig = field(default_factory=DataLoaderConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
    production: ProductionConfig = field(default_factory=ProductionConfig)
    checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
    loss: LossConfig = field(default_factory=LossConfig)
    tensorboard: TensorboardConfig = field(default_factory=TensorboardConfig)
    preprocessing: PreprocessingConfig = field(default_factory=PreprocessingConfig)
    # Class information (saved for inference)
    class_map: Optional[Dict[str, int]] = field(default=None, repr=False)
    class_names: Optional[List[str]] = field(default=None, repr=False)
    num_classes: Optional[int] = field(default=None, repr=False)  # Foreground classes only (model output = num_classes + 1 for background)

    def to_dict(self) -> Dict[str, Any]:
        """Convert config to dictionary, including class information."""
        config_dict = asdict(self)
        _normalize_legacy_loss_type(config_dict)
        # Convert Path objects to strings
        return self._convert_paths_to_strings(config_dict)

    def _convert_paths_to_strings(self, obj: Any) -> Any:
        """Recursively convert Path objects to strings."""
        if isinstance(obj, Path):
            return str(obj)
        elif isinstance(obj, dict):
            return {k: self._convert_paths_to_strings(v) for k, v in obj.items()}
        elif isinstance(obj, (list, tuple)):
            return [self._convert_paths_to_strings(item) for item in obj]
        else:
            return obj

    def save(self, path: Path) -> None:
        """Save configuration to JSON file.

        Args:
            path: Path to save the configuration file
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        config_dict = self.to_dict()

        with path.open("w", encoding="utf-8") as f:
            json.dump(config_dict, f, indent=2, default=str)

    @classmethod
    def load(cls, path: Path) -> "TrainingExperimentConfig":
        """Load configuration from JSON file with support for _base_ inheritance.

        Supports MMRotate-style nested config inheritance. Base configs are loaded
        and merged before converting to dataclass instances.

        Args:
            path: Path to the configuration file (can use _base_ field)

        Returns:
            TrainingExperimentConfig instance

        Examples:
            >>> # Simple config
            >>> config = TrainingExperimentConfig.load(Path("config.json"))

            >>> # Config with base inheritance
            >>> # config.json: {"_base_": "../_base_/datasets/dota.json", ...}
            >>> config = TrainingExperimentConfig.load(Path("config.json"))
        """
        path = Path(path)

        # Use enhanced config loader that supports _base_ inheritance
        frozen_cfg = load_config(path)
        config_dict = frozen_cfg.to_dict()
        _normalize_legacy_loss_type(config_dict)
        _normalize_legacy_model_keys(config_dict)

        # JSON `null` for a section should fall back to dataclass defaults, not be passed
        # as `None` into non-optional sub-configs.
        for _k in (
            "augmentation", "data_loader", "model", "training", "evaluation",
            "production", "checkpoint", "loss", "tensorboard", "preprocessing", "dataset",
        ):
            if _k in config_dict and config_dict[_k] is None:
                del config_dict[_k]

        valid_root = {f for f in cls.__dataclass_fields__}
        root_extra = set(config_dict) - valid_root
        if root_extra:
            raise ValueError(
                f"Unknown key(s) at config root: {sorted(root_extra)}. "
                f"Valid top-level keys: {sorted(valid_root)}"
            )

        # Convert nested dicts to dataclass instances (strict: no unknown keys per section)
        if "dataset" in config_dict and config_dict["dataset"] is not None:
            config_dict["dataset"] = DatasetConfig(
                **_strict_section(config_dict["dataset"], "dataset", DatasetConfig)
            )
        if "augmentation" in config_dict and config_dict["augmentation"] is not None:
            config_dict["augmentation"] = AugmentationConfig(
                **_strict_section(config_dict["augmentation"], "augmentation", AugmentationConfig)
            )
        if "data_loader" in config_dict and config_dict["data_loader"] is not None:
            config_dict["data_loader"] = DataLoaderConfig(
                **_strict_section(config_dict["data_loader"], "data_loader", DataLoaderConfig)
            )
        if "model" in config_dict and config_dict["model"] is not None:
            config_dict["model"] = ModelConfig(
                **_strict_section(config_dict["model"], "model", ModelConfig)
            )
        if "training" in config_dict and config_dict["training"] is not None:
            config_dict["training"] = TrainingConfig(
                **_strict_section(config_dict["training"], "training", TrainingConfig)
            )
        if "evaluation" in config_dict and config_dict["evaluation"] is not None:
            config_dict["evaluation"] = EvaluationConfig(
                **_strict_section(config_dict["evaluation"], "evaluation", EvaluationConfig)
            )
        if "production" in config_dict and config_dict["production"] is not None:
            config_dict["production"] = ProductionConfig(
                **_strict_section(config_dict["production"], "production", ProductionConfig)
            )
        if "checkpoint" in config_dict and config_dict["checkpoint"] is not None:
            config_dict["checkpoint"] = CheckpointConfig(
                **_strict_section(config_dict["checkpoint"], "checkpoint", CheckpointConfig)
            )
        if "loss" in config_dict and config_dict["loss"] is not None:
            config_dict["loss"] = LossConfig(
                **_strict_section(config_dict["loss"], "loss", LossConfig)
            )
        if "tensorboard" in config_dict and config_dict["tensorboard"] is not None:
            config_dict["tensorboard"] = TensorboardConfig(
                **_strict_section(config_dict["tensorboard"], "tensorboard", TensorboardConfig)
            )
        if "preprocessing" in config_dict and config_dict["preprocessing"] is not None:
            config_dict["preprocessing"] = PreprocessingConfig(
                **_strict_section(config_dict["preprocessing"], "preprocessing", PreprocessingConfig)
            )

        return cls(**config_dict)

    def print_summary(self) -> None:
        """Print a summary of the configuration."""
        print("Training Configuration Summary:")
        print("=" * 60)
        print(f"Model Type: {self.model_type}")
        if self.experiment_timestamp:
            print(f"Experiment Timestamp: {self.experiment_timestamp}")
        if self.git_commit:
            dirty = " dirty" if self.git_dirty else ""
            label = self.git_describe or self.git_commit[:12]
            print(f"Git: {label}{dirty}")
            if self.git_branch:
                print(f"Git branch: {self.git_branch}")
            if self.git_commit_date:
                print(f"Git commit date: {self.git_commit_date}")
        if self.package_version:
            print(f"Package: oriented-det {self.package_version}")
        print()

        if self.dataset:
            print("Dataset:")
            print(f"  Data Root: {self.dataset.data_root}")
            print(f"  Format: {self.dataset.format}")
            if self.dataset.format == "airbus_playground":
                print(f"  Annotations File: {self.dataset.annotations_file}")
                print(f"  Split File: {self.dataset.split_file}")
                print(f"  Val Split Id: {self.dataset.val_split_id}")
                if getattr(self.dataset, "train_includes_val", False):
                    print("  Train Includes Val: true (all folds in train; val fold for monitoring only)")
                print(f"  Ignore Labels: {self.dataset.ignore_labels}")
                print(f"  Map Labels: {self.dataset.map_labels}")
            else:
                if self.dataset.train_tiles_dirs:
                    print(f"  Train Tiles: {self.dataset.train_tiles_dirs}")
                else:
                    print(f"  Train Tiles: {self.dataset.train_tiles_dir}")
                if self.dataset.val_tiles_dirs:
                    print(f"  Val Tiles: {self.dataset.val_tiles_dirs}")
                else:
                    print(f"  Val Tiles: {self.dataset.val_tiles_dir}")
                if getattr(self.dataset, "same_folder", False):
                    print(f"  Same folder: images and labels in train/val tiles dirs")
            print(f"  difficult_strategy: {getattr(self.dataset, 'difficult_strategy', 'drop')}")
            print(f"  filter_empty_gt: {getattr(self.dataset, 'filter_empty_gt', False)}")
            print(f"  Tile overlap (px): {getattr(self.dataset, 'overlap', 16)}")
            print()

        print("DataLoader:")
        print(f"  Batch Size: {self.data_loader.batch_size}")
        print(f"  Num Workers: {self.data_loader.num_workers}")
        print(f"  Shuffle: {self.data_loader.shuffle}")
        print()

        print("Model:")
        print(f"  Backbone: {self.model.backbone}")
        print(f"  Pretrained: {self.model.pretrained_backbone}")
        print(f"  Anchor Scales: {self.model.anchor_scales}")
        print(f"  Anchor Ratios: {self.model.anchor_ratios}")
        print(f"  Inference Pre-NMS Score Threshold: {self.model.inference_pre_nms_score_threshold}")
        print()

        print("Training:")
        print(f"  Epochs: {self.training.num_epochs}")
        print(f"  Learning Rate: {self.training.learning_rate}")
        print(f"  Use LR Param Groups: {self.training.use_lr_param_groups}")
        if self.training.use_lr_param_groups:
            lh = getattr(self.training, "lr_mult_head", None)
            head_note = f", head(optional)={lh}" if lh is not None else ""
            print(
                "  LR Multipliers "
                f"(backbone/rpn/roi/other{head_note}): "
                f"{self.training.lr_mult_backbone}/"
                f"{self.training.lr_mult_rpn}/"
                f"{self.training.lr_mult_roi}/"
                f"{self.training.lr_mult_other}"
            )
        print(f"  Batch Size: {self.data_loader.batch_size}")
        print(f"  Gradient Accumulation: {self.training.gradient_accumulation_steps}")
        print(f"  Effective Batch Size: {self.data_loader.batch_size * self.training.gradient_accumulation_steps}")
        print(f"  Mixed Precision: {self.training.use_amp}")
        print()

        print("Evaluation:")
        print(f"  Score Threshold: {self.evaluation.score_threshold}")
        if self.evaluation.per_class_score_threshold:
            print(f"  Per-class score thresholds: {self.evaluation.per_class_score_threshold}")
        print(f"  Extended GT metrics: {self.evaluation.extended_gt_metrics}")
        print(f"  IoU Threshold: {self.evaluation.iou_threshold}")
        print(f"  Compute mAP final (best model): {self.evaluation.compute_map_final}")
        print(f"  Compute mAP every N epochs: {self.evaluation.compute_map_every_n_epochs}")
        print(f"  Use exact rotated IoU (mAP / GT cover): {self.evaluation.use_exact_rotated_iou}")
        _final_iou = getattr(self.evaluation, "use_exact_rotated_iou_for_final_map", None)
        if _final_iou is not None:
            print(f"  Use exact rotated IoU (final mAP): {_final_iou}")
        print()

        print("Production (overrides; null = use evaluation/model/dataset):")
        inf = self.production
        ow_px = resolve_inference_sliding_window_overlap_pixels(self)
        _eff_sc, _eff_pc, _eff_iou = effective_eval_metric_thresholds(self)
        print(f"  score_threshold: {inf.score_threshold!r} → effective {_eff_sc}")
        print(f"  per_class_score_threshold: {inf.per_class_score_threshold!r} → effective {_eff_pc!r}")
        print(f"  (mAP IoU from evaluation.iou_threshold: {self.evaluation.iou_threshold} → effective {_eff_iou})")
        print(f"  final_nms_iou_threshold: {inf.final_nms_iou_threshold!r}")
        print(f"  final_nms_use_cpu: {inf.final_nms_use_cpu!r}")
        print(f"  inference_pre_nms_score_threshold: {inf.inference_pre_nms_score_threshold!r}")
        print(f"  max_detections_per_image: {inf.max_detections_per_image!r}")
        print(f"  nms_class_agnostic: {inf.nms_class_agnostic!r}")
        print(f"  roi_inference_top_class_only: {inf.roi_inference_top_class_only!r}")
        print(f"  rpn_pre_nms_top_n: {inf.rpn_pre_nms_top_n!r}")
        print(f"  rpn_post_nms_top_n: {inf.rpn_post_nms_top_n!r}")
        print(f"  rpn_nms_threshold: {inf.rpn_nms_threshold!r}")
        print(f"  sliding overlap (px): config={inf.overlap_pixels!r} → effective {ow_px}")
        print(f"  ignore_margin_pixels: {inf.ignore_margin_pixels!r}")
        print(f"  use_first_image_canvas: {inf.use_first_image_canvas!r}")
        print(f"  stick_to_model_canvas: {inf.stick_to_model_canvas!r}")
        print()

        print("Checkpoint:")
        print(f"  Discover previous run: {getattr(self.checkpoint, 'discover_previous_run', False)}")
        print(f"  Resume from checkpoint epoch: {self.checkpoint.resume_from_checkpoint_epoch}")
        print(f"  Load Optimizer State: {self.checkpoint.load_optimizer_state}")
        print(f"  Include Prefixes: {self.checkpoint.load_include_prefixes}")
        print(f"  Exclude Prefixes: {self.checkpoint.load_exclude_prefixes}")
        print()

        print("Loss:")
        print(f"  Type: {self.loss.loss_type}")
        print(f"  Class Weight Method: {self.loss.class_weight_method}")
        print(f"  Background Weight: {self.loss.background_weight}")
        print()

        if self.num_classes is not None:
            print("Classes:")
            print(f"  Number of classes (foreground): {self.num_classes}")
            if self.class_names:
                print(f"  Class names ({len(self.class_names)}): {', '.join(self.class_names)}")
            if self.class_map:
                print(f"  Class map: {self.class_map}")

load(path) classmethod

Load configuration from JSON file with support for base inheritance.

Supports MMRotate-style nested config inheritance. Base configs are loaded and merged before converting to dataclass instances.

Parameters:

Name	Type	Description	Default
path	Path	Path to the configuration file (can use base field)	required

Returns:

Type	Description
TrainingExperimentConfig	TrainingExperimentConfig instance

Examples:

>>> # Simple config
>>> config = TrainingExperimentConfig.load(Path("config.json"))

>>> # Config with base inheritance
>>> # config.json: {"_base_": "../_base_/datasets/dota.json", ...}
>>> config = TrainingExperimentConfig.load(Path("config.json"))

Source code in oriented_det/train/config.py

@classmethod
def load(cls, path: Path) -> "TrainingExperimentConfig":
    """Load configuration from JSON file with support for _base_ inheritance.

    Supports MMRotate-style nested config inheritance. Base configs are loaded
    and merged before converting to dataclass instances.

    Args:
        path: Path to the configuration file (can use _base_ field)

    Returns:
        TrainingExperimentConfig instance

    Examples:
        >>> # Simple config
        >>> config = TrainingExperimentConfig.load(Path("config.json"))

        >>> # Config with base inheritance
        >>> # config.json: {"_base_": "../_base_/datasets/dota.json", ...}
        >>> config = TrainingExperimentConfig.load(Path("config.json"))
    """
    path = Path(path)

    # Use enhanced config loader that supports _base_ inheritance
    frozen_cfg = load_config(path)
    config_dict = frozen_cfg.to_dict()
    _normalize_legacy_loss_type(config_dict)
    _normalize_legacy_model_keys(config_dict)

    # JSON `null` for a section should fall back to dataclass defaults, not be passed
    # as `None` into non-optional sub-configs.
    for _k in (
        "augmentation", "data_loader", "model", "training", "evaluation",
        "production", "checkpoint", "loss", "tensorboard", "preprocessing", "dataset",
    ):
        if _k in config_dict and config_dict[_k] is None:
            del config_dict[_k]

    valid_root = {f for f in cls.__dataclass_fields__}
    root_extra = set(config_dict) - valid_root
    if root_extra:
        raise ValueError(
            f"Unknown key(s) at config root: {sorted(root_extra)}. "
            f"Valid top-level keys: {sorted(valid_root)}"
        )

    # Convert nested dicts to dataclass instances (strict: no unknown keys per section)
    if "dataset" in config_dict and config_dict["dataset"] is not None:
        config_dict["dataset"] = DatasetConfig(
            **_strict_section(config_dict["dataset"], "dataset", DatasetConfig)
        )
    if "augmentation" in config_dict and config_dict["augmentation"] is not None:
        config_dict["augmentation"] = AugmentationConfig(
            **_strict_section(config_dict["augmentation"], "augmentation", AugmentationConfig)
        )
    if "data_loader" in config_dict and config_dict["data_loader"] is not None:
        config_dict["data_loader"] = DataLoaderConfig(
            **_strict_section(config_dict["data_loader"], "data_loader", DataLoaderConfig)
        )
    if "model" in config_dict and config_dict["model"] is not None:
        config_dict["model"] = ModelConfig(
            **_strict_section(config_dict["model"], "model", ModelConfig)
        )
    if "training" in config_dict and config_dict["training"] is not None:
        config_dict["training"] = TrainingConfig(
            **_strict_section(config_dict["training"], "training", TrainingConfig)
        )
    if "evaluation" in config_dict and config_dict["evaluation"] is not None:
        config_dict["evaluation"] = EvaluationConfig(
            **_strict_section(config_dict["evaluation"], "evaluation", EvaluationConfig)
        )
    if "production" in config_dict and config_dict["production"] is not None:
        config_dict["production"] = ProductionConfig(
            **_strict_section(config_dict["production"], "production", ProductionConfig)
        )
    if "checkpoint" in config_dict and config_dict["checkpoint"] is not None:
        config_dict["checkpoint"] = CheckpointConfig(
            **_strict_section(config_dict["checkpoint"], "checkpoint", CheckpointConfig)
        )
    if "loss" in config_dict and config_dict["loss"] is not None:
        config_dict["loss"] = LossConfig(
            **_strict_section(config_dict["loss"], "loss", LossConfig)
        )
    if "tensorboard" in config_dict and config_dict["tensorboard"] is not None:
        config_dict["tensorboard"] = TensorboardConfig(
            **_strict_section(config_dict["tensorboard"], "tensorboard", TensorboardConfig)
        )
    if "preprocessing" in config_dict and config_dict["preprocessing"] is not None:
        config_dict["preprocessing"] = PreprocessingConfig(
            **_strict_section(config_dict["preprocessing"], "preprocessing", PreprocessingConfig)
        )

    return cls(**config_dict)

print_summary()

Print a summary of the configuration.

Source code in oriented_det/train/config.py

def print_summary(self) -> None:
    """Print a summary of the configuration."""
    print("Training Configuration Summary:")
    print("=" * 60)
    print(f"Model Type: {self.model_type}")
    if self.experiment_timestamp:
        print(f"Experiment Timestamp: {self.experiment_timestamp}")
    if self.git_commit:
        dirty = " dirty" if self.git_dirty else ""
        label = self.git_describe or self.git_commit[:12]
        print(f"Git: {label}{dirty}")
        if self.git_branch:
            print(f"Git branch: {self.git_branch}")
        if self.git_commit_date:
            print(f"Git commit date: {self.git_commit_date}")
    if self.package_version:
        print(f"Package: oriented-det {self.package_version}")
    print()

    if self.dataset:
        print("Dataset:")
        print(f"  Data Root: {self.dataset.data_root}")
        print(f"  Format: {self.dataset.format}")
        if self.dataset.format == "airbus_playground":
            print(f"  Annotations File: {self.dataset.annotations_file}")
            print(f"  Split File: {self.dataset.split_file}")
            print(f"  Val Split Id: {self.dataset.val_split_id}")
            if getattr(self.dataset, "train_includes_val", False):
                print("  Train Includes Val: true (all folds in train; val fold for monitoring only)")
            print(f"  Ignore Labels: {self.dataset.ignore_labels}")
            print(f"  Map Labels: {self.dataset.map_labels}")
        else:
            if self.dataset.train_tiles_dirs:
                print(f"  Train Tiles: {self.dataset.train_tiles_dirs}")
            else:
                print(f"  Train Tiles: {self.dataset.train_tiles_dir}")
            if self.dataset.val_tiles_dirs:
                print(f"  Val Tiles: {self.dataset.val_tiles_dirs}")
            else:
                print(f"  Val Tiles: {self.dataset.val_tiles_dir}")
            if getattr(self.dataset, "same_folder", False):
                print(f"  Same folder: images and labels in train/val tiles dirs")
        print(f"  difficult_strategy: {getattr(self.dataset, 'difficult_strategy', 'drop')}")
        print(f"  filter_empty_gt: {getattr(self.dataset, 'filter_empty_gt', False)}")
        print(f"  Tile overlap (px): {getattr(self.dataset, 'overlap', 16)}")
        print()

    print("DataLoader:")
    print(f"  Batch Size: {self.data_loader.batch_size}")
    print(f"  Num Workers: {self.data_loader.num_workers}")
    print(f"  Shuffle: {self.data_loader.shuffle}")
    print()

    print("Model:")
    print(f"  Backbone: {self.model.backbone}")
    print(f"  Pretrained: {self.model.pretrained_backbone}")
    print(f"  Anchor Scales: {self.model.anchor_scales}")
    print(f"  Anchor Ratios: {self.model.anchor_ratios}")
    print(f"  Inference Pre-NMS Score Threshold: {self.model.inference_pre_nms_score_threshold}")
    print()

    print("Training:")
    print(f"  Epochs: {self.training.num_epochs}")
    print(f"  Learning Rate: {self.training.learning_rate}")
    print(f"  Use LR Param Groups: {self.training.use_lr_param_groups}")
    if self.training.use_lr_param_groups:
        lh = getattr(self.training, "lr_mult_head", None)
        head_note = f", head(optional)={lh}" if lh is not None else ""
        print(
            "  LR Multipliers "
            f"(backbone/rpn/roi/other{head_note}): "
            f"{self.training.lr_mult_backbone}/"
            f"{self.training.lr_mult_rpn}/"
            f"{self.training.lr_mult_roi}/"
            f"{self.training.lr_mult_other}"
        )
    print(f"  Batch Size: {self.data_loader.batch_size}")
    print(f"  Gradient Accumulation: {self.training.gradient_accumulation_steps}")
    print(f"  Effective Batch Size: {self.data_loader.batch_size * self.training.gradient_accumulation_steps}")
    print(f"  Mixed Precision: {self.training.use_amp}")
    print()

    print("Evaluation:")
    print(f"  Score Threshold: {self.evaluation.score_threshold}")
    if self.evaluation.per_class_score_threshold:
        print(f"  Per-class score thresholds: {self.evaluation.per_class_score_threshold}")
    print(f"  Extended GT metrics: {self.evaluation.extended_gt_metrics}")
    print(f"  IoU Threshold: {self.evaluation.iou_threshold}")
    print(f"  Compute mAP final (best model): {self.evaluation.compute_map_final}")
    print(f"  Compute mAP every N epochs: {self.evaluation.compute_map_every_n_epochs}")
    print(f"  Use exact rotated IoU (mAP / GT cover): {self.evaluation.use_exact_rotated_iou}")
    _final_iou = getattr(self.evaluation, "use_exact_rotated_iou_for_final_map", None)
    if _final_iou is not None:
        print(f"  Use exact rotated IoU (final mAP): {_final_iou}")
    print()

    print("Production (overrides; null = use evaluation/model/dataset):")
    inf = self.production
    ow_px = resolve_inference_sliding_window_overlap_pixels(self)
    _eff_sc, _eff_pc, _eff_iou = effective_eval_metric_thresholds(self)
    print(f"  score_threshold: {inf.score_threshold!r} → effective {_eff_sc}")
    print(f"  per_class_score_threshold: {inf.per_class_score_threshold!r} → effective {_eff_pc!r}")
    print(f"  (mAP IoU from evaluation.iou_threshold: {self.evaluation.iou_threshold} → effective {_eff_iou})")
    print(f"  final_nms_iou_threshold: {inf.final_nms_iou_threshold!r}")
    print(f"  final_nms_use_cpu: {inf.final_nms_use_cpu!r}")
    print(f"  inference_pre_nms_score_threshold: {inf.inference_pre_nms_score_threshold!r}")
    print(f"  max_detections_per_image: {inf.max_detections_per_image!r}")
    print(f"  nms_class_agnostic: {inf.nms_class_agnostic!r}")
    print(f"  roi_inference_top_class_only: {inf.roi_inference_top_class_only!r}")
    print(f"  rpn_pre_nms_top_n: {inf.rpn_pre_nms_top_n!r}")
    print(f"  rpn_post_nms_top_n: {inf.rpn_post_nms_top_n!r}")
    print(f"  rpn_nms_threshold: {inf.rpn_nms_threshold!r}")
    print(f"  sliding overlap (px): config={inf.overlap_pixels!r} → effective {ow_px}")
    print(f"  ignore_margin_pixels: {inf.ignore_margin_pixels!r}")
    print(f"  use_first_image_canvas: {inf.use_first_image_canvas!r}")
    print(f"  stick_to_model_canvas: {inf.stick_to_model_canvas!r}")
    print()

    print("Checkpoint:")
    print(f"  Discover previous run: {getattr(self.checkpoint, 'discover_previous_run', False)}")
    print(f"  Resume from checkpoint epoch: {self.checkpoint.resume_from_checkpoint_epoch}")
    print(f"  Load Optimizer State: {self.checkpoint.load_optimizer_state}")
    print(f"  Include Prefixes: {self.checkpoint.load_include_prefixes}")
    print(f"  Exclude Prefixes: {self.checkpoint.load_exclude_prefixes}")
    print()

    print("Loss:")
    print(f"  Type: {self.loss.loss_type}")
    print(f"  Class Weight Method: {self.loss.class_weight_method}")
    print(f"  Background Weight: {self.loss.background_weight}")
    print()

    if self.num_classes is not None:
        print("Classes:")
        print(f"  Number of classes (foreground): {self.num_classes}")
        if self.class_names:
            print(f"  Class names ({len(self.class_names)}): {', '.join(self.class_names)}")
        if self.class_map:
            print(f"  Class map: {self.class_map}")

save(path)

Save configuration to JSON file.

Parameters:

Name	Type	Description	Default
path	Path	Path to save the configuration file	required

Source code in oriented_det/train/config.py

def save(self, path: Path) -> None:
    """Save configuration to JSON file.

    Args:
        path: Path to save the configuration file
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    config_dict = self.to_dict()

    with path.open("w", encoding="utf-8") as f:
        json.dump(config_dict, f, indent=2, default=str)

to_dict()

Convert config to dictionary, including class information.

Source code in oriented_det/train/config.py

def to_dict(self) -> Dict[str, Any]:
    """Convert config to dictionary, including class information."""
    config_dict = asdict(self)
    _normalize_legacy_loss_type(config_dict)
    # Convert Path objects to strings
    return self._convert_paths_to_strings(config_dict)

TrainingProfiler

Context manager for profiling PyTorch training.

This profiler helps identify bottlenecks in training by tracking: - GPU kernel execution time - CPU operations - Memory usage - Data loading time - Model forward/backward pass time

Example

from oriented_det.train.profiler import TrainingProfiler

profiler = TrainingProfiler(
    log_dir="runs/profiling",
    activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
    profile_memory=True,
)

with profiler:
    # Your training code here
    for batch in train_loader:
        loss = model(batch)
        loss.backward()
        optimizer.step()

# View results in TensorBoard:
# tensorboard --logdir runs/profiling

Source code in oriented_det/train/profiler.py

class TrainingProfiler:
    """Context manager for profiling PyTorch training.

    This profiler helps identify bottlenecks in training by tracking:
    - GPU kernel execution time
    - CPU operations
    - Memory usage
    - Data loading time
    - Model forward/backward pass time

    Example:
        ```python
        from oriented_det.train.profiler import TrainingProfiler

        profiler = TrainingProfiler(
            log_dir="runs/profiling",
            activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
            profile_memory=True,
        )

        with profiler:
            # Your training code here
            for batch in train_loader:
                loss = model(batch)
                loss.backward()
                optimizer.step()

        # View results in TensorBoard:
        # tensorboard --logdir runs/profiling
        ```
    """

    def __init__(
        self,
        log_dir: str | Path,
        activities: Optional[list] = None,
        schedule: Optional[Any] = None,
        on_trace_ready: Optional[Any] = None,
        record_shapes: bool = False,
        profile_memory: bool = True,
        with_stack: bool = False,
        with_flops: bool = False,
        experimental_config: Optional[Any] = None,
    ):
        """Initialize training profiler.

        Args:
            log_dir: Directory to save profiling traces (for TensorBoard)
            activities: List of activities to profile (default: [CUDA, CPU])
            schedule: Profiling schedule (default: warmup 1, active 2, repeat 1)
            on_trace_ready: Callback when trace is ready (default: TensorBoard handler)
            record_shapes: Whether to record tensor shapes
            profile_memory: Whether to profile memory usage
            with_stack: Whether to record stack traces
            with_flops: Whether to estimate FLOPs (requires torch >= 1.12)
            experimental_config: Experimental profiler configuration
        """
        _require_torch()

        if activities is None:
            activities = []
            if torch.cuda.is_available():
                activities.append(ProfilerActivity.CUDA)
            activities.append(ProfilerActivity.CPU)

        if schedule is None and profiler_schedule is not None:
            # Default schedule: warmup 1 step, active 2 steps, repeat 1 time
            schedule = profiler_schedule(
                wait=1,  # Skip first step (warmup)
                warmup=1,  # Warmup for 1 step
                active=2,  # Profile for 2 steps
                repeat=1,  # Repeat 1 time
            )

        if on_trace_ready is None:
            log_dir = Path(log_dir)
            log_dir.mkdir(parents=True, exist_ok=True)
            on_trace_ready = tensorboard_trace_handler(str(log_dir))

        self.log_dir = Path(log_dir)
        self.profiler = profile(
            activities=activities,
            schedule=schedule,
            on_trace_ready=on_trace_ready,
            record_shapes=record_shapes,
            profile_memory=profile_memory,
            with_stack=with_stack,
            with_flops=with_flops,
            experimental_config=experimental_config,
        )

    def __enter__(self):
        """Enter profiling context."""
        self.profiler.__enter__()
        return self

    def __exit__(self, *args):
        """Exit profiling context."""
        return self.profiler.__exit__(*args)

    def step(self):
        """Advance profiler to next step."""
        self.profiler.step()

    def export_chrome_trace(self, path: str | Path):
        """Export trace to Chrome trace format.

        Args:
            path: Path to save trace file (.json)
        """
        self.profiler.export_chrome_trace(str(path))

    def export_stacks(self, path: str | Path, metric: str = "self_cuda_time_total"):
        """Export stack traces.

        Args:
            path: Path to save stacks file (.txt)
            metric: Metric to sort by (e.g., "self_cuda_time_total", "self_cpu_time_total")
        """
        self.profiler.export_stacks(str(path), metric=metric)

    def key_averages(self, group_by_input_shapes: bool = False, group_by_stack_n: int = 0):
        """Get key averages from profiler.

        Args:
            group_by_input_shapes: Group by input tensor shapes
            group_by_stack_n: Group by stack trace (n levels)

        Returns:
            EventList with averaged events
        """
        return self.profiler.key_averages(
            group_by_input_shapes=group_by_input_shapes,
            group_by_stack_n=group_by_stack_n,
        )

    def print_summary(self, sort_by: str = "cuda_time_total", row_limit: int = 20):
        """Print profiling summary.

        Args:
            sort_by: Metric to sort by (e.g., "cuda_time_total", "cpu_time_total", "self_cuda_time_total")
            row_limit: Maximum number of rows to print
        """
        events = self.key_averages()
        print(f"\n{'='*80}")
        print(f"Profiling Summary (sorted by {sort_by})")
        print(f"{'='*80}")
        print(events.table(sort_by=sort_by, row_limit=row_limit))
        print(f"{'='*80}\n")

__enter__()

Enter profiling context.

Source code in oriented_det/train/profiler.py

def __enter__(self):
    """Enter profiling context."""
    self.profiler.__enter__()
    return self

__exit__(*args)

Exit profiling context.

Source code in oriented_det/train/profiler.py

def __exit__(self, *args):
    """Exit profiling context."""
    return self.profiler.__exit__(*args)

__init__(log_dir, activities=None, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=True, with_stack=False, with_flops=False, experimental_config=None)

Initialize training profiler.

Parameters:

Name	Type	Description	Default
log_dir	str | Path	Directory to save profiling traces (for TensorBoard)	required
activities	Optional[list]	List of activities to profile (default: [CUDA, CPU])	None
schedule	Optional[Any]	Profiling schedule (default: warmup 1, active 2, repeat 1)	None
on_trace_ready	Optional[Any]	Callback when trace is ready (default: TensorBoard handler)	None
record_shapes	bool	Whether to record tensor shapes	False
profile_memory	bool	Whether to profile memory usage	True
with_stack	bool	Whether to record stack traces	False
with_flops	bool	Whether to estimate FLOPs (requires torch >= 1.12)	False
experimental_config	Optional[Any]	Experimental profiler configuration	None

Source code in oriented_det/train/profiler.py

def __init__(
    self,
    log_dir: str | Path,
    activities: Optional[list] = None,
    schedule: Optional[Any] = None,
    on_trace_ready: Optional[Any] = None,
    record_shapes: bool = False,
    profile_memory: bool = True,
    with_stack: bool = False,
    with_flops: bool = False,
    experimental_config: Optional[Any] = None,
):
    """Initialize training profiler.

    Args:
        log_dir: Directory to save profiling traces (for TensorBoard)
        activities: List of activities to profile (default: [CUDA, CPU])
        schedule: Profiling schedule (default: warmup 1, active 2, repeat 1)
        on_trace_ready: Callback when trace is ready (default: TensorBoard handler)
        record_shapes: Whether to record tensor shapes
        profile_memory: Whether to profile memory usage
        with_stack: Whether to record stack traces
        with_flops: Whether to estimate FLOPs (requires torch >= 1.12)
        experimental_config: Experimental profiler configuration
    """
    _require_torch()

    if activities is None:
        activities = []
        if torch.cuda.is_available():
            activities.append(ProfilerActivity.CUDA)
        activities.append(ProfilerActivity.CPU)

    if schedule is None and profiler_schedule is not None:
        # Default schedule: warmup 1 step, active 2 steps, repeat 1 time
        schedule = profiler_schedule(
            wait=1,  # Skip first step (warmup)
            warmup=1,  # Warmup for 1 step
            active=2,  # Profile for 2 steps
            repeat=1,  # Repeat 1 time
        )

    if on_trace_ready is None:
        log_dir = Path(log_dir)
        log_dir.mkdir(parents=True, exist_ok=True)
        on_trace_ready = tensorboard_trace_handler(str(log_dir))

    self.log_dir = Path(log_dir)
    self.profiler = profile(
        activities=activities,
        schedule=schedule,
        on_trace_ready=on_trace_ready,
        record_shapes=record_shapes,
        profile_memory=profile_memory,
        with_stack=with_stack,
        with_flops=with_flops,
        experimental_config=experimental_config,
    )

export_chrome_trace(path)

Export trace to Chrome trace format.

Parameters:

Name	Type	Description	Default
path	str | Path	Path to save trace file (.json)	required

Source code in oriented_det/train/profiler.py

def export_chrome_trace(self, path: str | Path):
    """Export trace to Chrome trace format.

    Args:
        path: Path to save trace file (.json)
    """
    self.profiler.export_chrome_trace(str(path))

export_stacks(path, metric='self_cuda_time_total')

Export stack traces.

Parameters:

Name	Type	Description	Default
path	str | Path	Path to save stacks file (.txt)	required
metric	str	Metric to sort by (e.g., "self_cuda_time_total", "self_cpu_time_total")	'self_cuda_time_total'

Source code in oriented_det/train/profiler.py

def export_stacks(self, path: str | Path, metric: str = "self_cuda_time_total"):
    """Export stack traces.

    Args:
        path: Path to save stacks file (.txt)
        metric: Metric to sort by (e.g., "self_cuda_time_total", "self_cpu_time_total")
    """
    self.profiler.export_stacks(str(path), metric=metric)

key_averages(group_by_input_shapes=False, group_by_stack_n=0)

Get key averages from profiler.

Parameters:

Name	Type	Description	Default
group_by_input_shapes	bool	Group by input tensor shapes	False
group_by_stack_n	int	Group by stack trace (n levels)	0

Returns:

Type	Description
	EventList with averaged events

Source code in oriented_det/train/profiler.py

def key_averages(self, group_by_input_shapes: bool = False, group_by_stack_n: int = 0):
    """Get key averages from profiler.

    Args:
        group_by_input_shapes: Group by input tensor shapes
        group_by_stack_n: Group by stack trace (n levels)

    Returns:
        EventList with averaged events
    """
    return self.profiler.key_averages(
        group_by_input_shapes=group_by_input_shapes,
        group_by_stack_n=group_by_stack_n,
    )

print_summary(sort_by='cuda_time_total', row_limit=20)

Print profiling summary.

Parameters:

Name	Type	Description	Default
sort_by	str	Metric to sort by (e.g., "cuda_time_total", "cpu_time_total", "self_cuda_time_total")	'cuda_time_total'
row_limit	int	Maximum number of rows to print	20

Source code in oriented_det/train/profiler.py

def print_summary(self, sort_by: str = "cuda_time_total", row_limit: int = 20):
    """Print profiling summary.

    Args:
        sort_by: Metric to sort by (e.g., "cuda_time_total", "cpu_time_total", "self_cuda_time_total")
        row_limit: Maximum number of rows to print
    """
    events = self.key_averages()
    print(f"\n{'='*80}")
    print(f"Profiling Summary (sorted by {sort_by})")
    print(f"{'='*80}")
    print(events.table(sort_by=sort_by, row_limit=row_limit))
    print(f"{'='*80}\n")

step()

Advance profiler to next step.

Source code in oriented_det/train/profiler.py

def step(self):
    """Advance profiler to next step."""
    self.profiler.step()

WarmupScheduler

Learning rate scheduler with warmup that wraps a base scheduler.

During warmup, the learning rate increases linearly from 0 to the base learning rate. After warmup, it uses the base scheduler (stepped per epoch, not per optimizer step).

This is useful for stabilizing training in the early stages, especially when using large learning rates or training from scratch.

Example

optimizer = optim.SGD(model.parameters(), lr=0.001) base_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.1) warmup_scheduler = WarmupScheduler(optimizer, base_scheduler, warmup_steps=500)

During training loop:

for epoch in range(num_epochs): for batch in dataloader: # ... training step ... optimizer.step() warmup_scheduler.step() # Step per optimizer step during warmup warmup_scheduler.step_epoch() # Step base scheduler per epoch

Source code in oriented_det/train/utils.py

class WarmupScheduler:
    """Learning rate scheduler with warmup that wraps a base scheduler.

    During warmup, the learning rate increases linearly from 0 to the base learning rate.
    After warmup, it uses the base scheduler (stepped per epoch, not per optimizer step).

    This is useful for stabilizing training in the early stages, especially when using
    large learning rates or training from scratch.

    Example:
        >>> optimizer = optim.SGD(model.parameters(), lr=0.001)
        >>> base_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.1)
        >>> warmup_scheduler = WarmupScheduler(optimizer, base_scheduler, warmup_steps=500)
        >>> 
        >>> # During training loop:
        >>> for epoch in range(num_epochs):
        >>>     for batch in dataloader:
        >>>         # ... training step ...
        >>>         optimizer.step()
        >>>         warmup_scheduler.step()  # Step per optimizer step during warmup
        >>>     warmup_scheduler.step_epoch()  # Step base scheduler per epoch
    """

    def __init__(self, optimizer: optim.Optimizer, base_scheduler: optim.lr_scheduler._LRScheduler, warmup_steps: int):
        """Initialize warmup scheduler.

        Args:
            optimizer: The optimizer whose learning rate will be scheduled
            base_scheduler: The base scheduler (e.g., StepLR) to use after warmup
            warmup_steps: Number of optimizer steps for warmup (not epochs)
        """
        _require_torch()
        if optim is None:
            raise RuntimeError("torch.optim is required for WarmupScheduler.")

        self.optimizer = optimizer
        self.base_scheduler = base_scheduler
        self.warmup_steps = warmup_steps
        self.base_lrs = [float(param_group['lr']) for param_group in optimizer.param_groups]
        self.current_step = 0
        self.warmup_complete = False

        # Set initial learning rate to 0
        for param_group in optimizer.param_groups:
            param_group['lr'] = 0.0

    def step(self) -> None:
        """Step the scheduler (called per optimizer step during warmup).

        This should be called after each optimizer.step() during the warmup period.
        After warmup completes, this method does nothing and step_epoch() should
        be used instead to step the base scheduler.
        """
        if not self.warmup_complete:
            self.current_step += 1

            if self.current_step <= self.warmup_steps:
                # Linear warmup per param group.
                progress = self.current_step / self.warmup_steps
                for base_lr, param_group in zip(self.base_lrs, self.optimizer.param_groups):
                    param_group['lr'] = base_lr * progress
            else:
                # Warmup complete - restore each group's base LR and mark as complete
                # The base scheduler will handle further updates per epoch
                for base_lr, param_group in zip(self.base_lrs, self.optimizer.param_groups):
                    param_group['lr'] = base_lr
                self.warmup_complete = True

    def step_epoch(self) -> None:
        """Step the base scheduler (called per epoch after warmup).

        This should be called once per epoch after warmup completes.
        During warmup, this method does nothing.
        """
        if self.warmup_complete:
            self.base_scheduler.step()

    def get_last_lr(self) -> List[float]:
        """Get the last learning rate for each parameter group."""
        return [param_group['lr'] for param_group in self.optimizer.param_groups]

    def state_dict(self) -> Dict[str, Any]:
        """Get state dict for checkpointing."""
        return {
            'current_step': self.current_step,
            'warmup_complete': self.warmup_complete,
            'base_scheduler': self.base_scheduler.state_dict(),
        }

    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
        """Load state dict from checkpoint."""
        self.current_step = state_dict['current_step']
        self.warmup_complete = state_dict.get('warmup_complete', False)
        self.base_scheduler.load_state_dict(state_dict['base_scheduler'])

__init__(optimizer, base_scheduler, warmup_steps)

Initialize warmup scheduler.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	The optimizer whose learning rate will be scheduled	required
base_scheduler	_LRScheduler	The base scheduler (e.g., StepLR) to use after warmup	required
warmup_steps	int	Number of optimizer steps for warmup (not epochs)	required

Source code in oriented_det/train/utils.py

def __init__(self, optimizer: optim.Optimizer, base_scheduler: optim.lr_scheduler._LRScheduler, warmup_steps: int):
    """Initialize warmup scheduler.

    Args:
        optimizer: The optimizer whose learning rate will be scheduled
        base_scheduler: The base scheduler (e.g., StepLR) to use after warmup
        warmup_steps: Number of optimizer steps for warmup (not epochs)
    """
    _require_torch()
    if optim is None:
        raise RuntimeError("torch.optim is required for WarmupScheduler.")

    self.optimizer = optimizer
    self.base_scheduler = base_scheduler
    self.warmup_steps = warmup_steps
    self.base_lrs = [float(param_group['lr']) for param_group in optimizer.param_groups]
    self.current_step = 0
    self.warmup_complete = False

    # Set initial learning rate to 0
    for param_group in optimizer.param_groups:
        param_group['lr'] = 0.0

get_last_lr()

Get the last learning rate for each parameter group.

Source code in oriented_det/train/utils.py

def get_last_lr(self) -> List[float]:
    """Get the last learning rate for each parameter group."""
    return [param_group['lr'] for param_group in self.optimizer.param_groups]

load_state_dict(state_dict)

Load state dict from checkpoint.

Source code in oriented_det/train/utils.py

def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
    """Load state dict from checkpoint."""
    self.current_step = state_dict['current_step']
    self.warmup_complete = state_dict.get('warmup_complete', False)
    self.base_scheduler.load_state_dict(state_dict['base_scheduler'])

state_dict()

Get state dict for checkpointing.

Source code in oriented_det/train/utils.py

def state_dict(self) -> Dict[str, Any]:
    """Get state dict for checkpointing."""
    return {
        'current_step': self.current_step,
        'warmup_complete': self.warmup_complete,
        'base_scheduler': self.base_scheduler.state_dict(),
    }

step()

Step the scheduler (called per optimizer step during warmup).

This should be called after each optimizer.step() during the warmup period. After warmup completes, this method does nothing and step_epoch() should be used instead to step the base scheduler.

Source code in oriented_det/train/utils.py

def step(self) -> None:
    """Step the scheduler (called per optimizer step during warmup).

    This should be called after each optimizer.step() during the warmup period.
    After warmup completes, this method does nothing and step_epoch() should
    be used instead to step the base scheduler.
    """
    if not self.warmup_complete:
        self.current_step += 1

        if self.current_step <= self.warmup_steps:
            # Linear warmup per param group.
            progress = self.current_step / self.warmup_steps
            for base_lr, param_group in zip(self.base_lrs, self.optimizer.param_groups):
                param_group['lr'] = base_lr * progress
        else:
            # Warmup complete - restore each group's base LR and mark as complete
            # The base scheduler will handle further updates per epoch
            for base_lr, param_group in zip(self.base_lrs, self.optimizer.param_groups):
                param_group['lr'] = base_lr
            self.warmup_complete = True

step_epoch()

Step the base scheduler (called per epoch after warmup).

This should be called once per epoch after warmup completes. During warmup, this method does nothing.

Source code in oriented_det/train/utils.py

def step_epoch(self) -> None:
    """Step the base scheduler (called per epoch after warmup).

    This should be called once per epoch after warmup completes.
    During warmup, this method does nothing.
    """
    if self.warmup_complete:
        self.base_scheduler.step()

apply_inference_config_to_model(model, inf)

Patch decode / NMS-related attributes on model from non-null production.* fields.

Call only for checkpoint-based inference outside the training loop — e.g. tools/save_predictions.load_model_from_checkpoint (deploy, image_demo, test_single_image). Not used by tools/train.py: during training the live module keeps model.* only so deploy-oriented RPN/NMS overrides cannot slow or alter the train forward pass.

Attributes that do not exist on model (e.g. RPN keys on RetinaNet) are skipped.

Source code in oriented_det/train/config.py

def apply_inference_config_to_model(model: Any, inf: Optional[ProductionConfig]) -> None:
    """Patch decode / NMS-related attributes on ``model`` from non-null ``production.*`` fields.

    Call only for **checkpoint-based inference** outside the training loop — e.g.
    ``tools/save_predictions.load_model_from_checkpoint`` (deploy, ``image_demo``, ``test_single_image``).
    **Not** used by ``tools/train.py``: during training the live module keeps ``model.*`` only so
    deploy-oriented RPN/NMS overrides cannot slow or alter the train forward pass.

    Attributes that do not exist on ``model`` (e.g. RPN keys on RetinaNet) are skipped.
    """
    if inf is None:
        return

    def _set(name: str, value: Any) -> None:
        if value is None:
            return
        if hasattr(model, name):
            setattr(model, name, value)

    _set("inference_pre_nms_score_threshold", inf.inference_pre_nms_score_threshold)
    if inf.inference_pre_nms_score_threshold is not None and not hasattr(
        model, "inference_pre_nms_score_threshold"
    ):
        _set("score_threshold", inf.inference_pre_nms_score_threshold)

    _set("final_nms_iou_threshold", inf.final_nms_iou_threshold)
    _set("final_nms_use_cpu", inf.final_nms_use_cpu)
    _set("max_detections_per_image", inf.max_detections_per_image)
    _set("nms_class_agnostic", inf.nms_class_agnostic)
    _set("roi_inference_top_class_only", inf.roi_inference_top_class_only)
    _set("rpn_pre_nms_top_n", inf.rpn_pre_nms_top_n)
    _set("rpn_post_nms_top_n", inf.rpn_post_nms_top_n)
    _set("rpn_nms_threshold", inf.rpn_nms_threshold)

build_grouped_ce_spec(groups, class_map, num_foreground_classes)

Map config group names to 1-indexed class id lists.

Each class name may appear in at most one group. Classes not listed in any group use fine-grained CE only (even when grouped_alpha > 0).

Source code in oriented_det/train/grouped_ce.py

def build_grouped_ce_spec(
    groups: Dict[str, Sequence[str]],
    class_map: Dict[str, int],
    num_foreground_classes: int,
) -> GroupedCeSpec:
    """Map config group names to 1-indexed class id lists.

    Each class name may appear in at most one group. Classes not listed in any group
    use fine-grained CE only (even when grouped_alpha > 0).
    """
    if num_foreground_classes < 1:
        raise ValueError("num_foreground_classes must be >= 1 for grouped CE")

    name_to_group: Dict[str, int] = {}
    group_lists: List[List[int]] = []
    for g_idx, class_names in enumerate(groups.values()):
        ids: List[int] = []
        for name in class_names:
            if name not in class_map:
                continue
            cid = int(class_map[name])
            if cid < 1 or cid > num_foreground_classes:
                continue
            if name in name_to_group:
                raise ValueError(
                    f"Class {name!r} appears in multiple roi_grouped_ce groups "
                    f"(groups {name_to_group[name]} and {g_idx})"
                )
            name_to_group[name] = g_idx
            ids.append(cid)
        group_lists.append(sorted(set(ids)))

    class_in_group = [-1] * (num_foreground_classes + 1)
    for g_idx, ids in enumerate(group_lists):
        for cid in ids:
            class_in_group[cid] = g_idx

    frozen_lists = tuple(tuple(ids) for ids in group_lists)
    return GroupedCeSpec(group_index_lists=frozen_lists, class_in_group_id=tuple(class_in_group))

collate_dota_samples(batch)

Collate function for DOTA dataset samples.

Converts DOTASample objects to format expected by models. Loads images from paths and converts them to tensors.

Source code in oriented_det/train/utils.py

def collate_dota_samples(batch: List[Any]) -> Tuple[List[Any], List[Dict[str, Any]]]:
    """Collate function for DOTA dataset samples.

    Converts DOTASample objects to format expected by models.
    Loads images from paths and converts them to tensors.
    """
    _require_torch()

    try:
        from PIL import Image
    except ImportError:
        raise RuntimeError("PIL/Pillow is required to load images.")

    try:
        from torchvision import transforms as T
    except ImportError:
        raise RuntimeError("torchvision is required for image transforms.")

    images = []
    targets = []

    # Convert PIL Image to tensor transform
    to_tensor = T.ToTensor()

    for sample in batch:
        # Load image from path
        image_path = sample.image_path
        try:
            pil_image = Image.open(image_path).convert("RGB")
            # Convert to tensor (C, H, W) format with values in [0, 1]
            image_tensor = to_tensor(pil_image)
            images.append(image_tensor)
        except Exception as e:
            raise RuntimeError(f"Failed to load image from {image_path}: {e}") from e

        # Convert annotations to target format
        # Apply le90 normalization (MMRotate standard for DOTA)
        # This ensures width >= height and angles in [-π/2, π/2) range
        from oriented_det.geometry.rbox import normalize_le90

        rboxes = []
        labels = []
        class_to_id = {}
        for ann in sample.annotations:
            if ann.class_name not in class_to_id:
                class_to_id[ann.class_name] = len(class_to_id)
            # Normalize to le90 convention (width >= height, angle in [-π/2, π/2))
            normalized_rbox = normalize_le90(ann.rbox)
            rboxes.append(normalized_rbox)
            labels.append(class_to_id[ann.class_name])

        target = {
            "rboxes": rboxes,
            "labels": torch.tensor(labels, dtype=torch.int64),
            "image_id": sample.image_path.stem if hasattr(sample, "image_path") else None,
            "image_filename": Path(sample.image_path).name
            if hasattr(sample, "image_path") and sample.image_path
            else None,
        }

        targets.append(target)

    return images, targets

collate_fn_generic(batch)

Generic collate function for detection datasets.

Expects batch items to be dicts with 'image' and 'target' keys.

Source code in oriented_det/train/utils.py

def collate_fn_generic(batch: List[Dict[str, Any]]) -> Tuple[List[Any], List[Dict[str, Any]]]:
    """Generic collate function for detection datasets.

    Expects batch items to be dicts with 'image' and 'target' keys.
    """
    _require_torch()

    images = []
    targets = []

    for item in batch:
        images.append(item.get("image"))
        targets.append(item.get("target", {}))

    return images, targets

configure_roi_grouped_ce(model, loss_config, class_map, *, num_foreground_classes, device=None)

Attach grouped CE schedule and group indices to a two-stage detector model.

Returns True when grouped CE is enabled and configured.

Source code in oriented_det/train/grouped_ce.py

def configure_roi_grouped_ce(
    model: object,
    loss_config: object,
    class_map: Optional[Dict[str, int]],
    *,
    num_foreground_classes: int,
    device: Optional["torch.device"] = None,
) -> bool:
    """Attach grouped CE schedule and group indices to a two-stage detector model.

    Returns True when grouped CE is enabled and configured.
    """
    if torch is None:
        raise RuntimeError("PyTorch is required for grouped CE configuration.")

    enabled = bool(getattr(loss_config, "roi_grouped_ce_enabled", False))
    groups = getattr(loss_config, "roi_grouped_ce_groups", None) or {}
    if not enabled or not groups:
        if hasattr(model, "clear_roi_grouped_ce"):
            model.clear_roi_grouped_ce()
        return False

    if not class_map:
        raise ValueError("roi_grouped_ce requires class_map from the training dataset")

    spec = build_grouped_ce_spec(groups, class_map, num_foreground_classes)
    if not any(spec.group_index_lists):
        raise ValueError("roi_grouped_ce_groups resolved to no class ids; check class names")

    if device is None:
        if hasattr(model, "parameters"):
            params = list(model.parameters())
            device = params[0].device if params else torch.device("cpu")
        else:
            device = torch.device("cpu")

    class_in_group_t = torch.tensor(spec.class_in_group_id, dtype=torch.long, device=device)

    if not hasattr(model, "set_roi_grouped_ce"):
        raise TypeError(f"{type(model).__name__} does not support ROI grouped CE")

    model.set_roi_grouped_ce(
        group_index_lists=list(spec.group_index_lists),
        class_in_group_id=class_in_group_t,
        schedule_type=getattr(loss_config, "roi_grouped_ce_schedule_type", None),
        schedule_start_epoch=int(getattr(loss_config, "roi_grouped_ce_schedule_start_epoch", 0) or 0),
        schedule_end_epoch=int(getattr(loss_config, "roi_grouped_ce_schedule_end_epoch", 0) or 0),
        schedule_power=float(getattr(loss_config, "roi_grouped_ce_schedule_power", 1.0) or 1.0),
    )
    return True

create_cosine_with_tail_lr_scheduler(optimizer, training, *, last_epoch=-1)

Cosine decay then constant lr_scheduler_cosine_tail_lr (no PyTorch restart).

Source code in oriented_det/train/utils.py

def create_cosine_with_tail_lr_scheduler(
    optimizer: "optim.Optimizer",
    training: Any,
    *,
    last_epoch: int = -1,
) -> Tuple[Any, int, int, float]:
    """Cosine decay then constant ``lr_scheduler_cosine_tail_lr`` (no PyTorch restart)."""
    _require_torch()
    eta_min = float(getattr(training, "lr_scheduler_cosine_eta_min", 1e-6))
    tail_lr_cfg = getattr(training, "lr_scheduler_cosine_tail_lr", None)
    tail_lr = float(eta_min if tail_lr_cfg is None else tail_lr_cfg)
    cosine_epochs, tail_epochs = resolve_cosine_with_tail_lengths(training)
    scheduler = CosineAnnealingWithFixedTailLR(
        optimizer,
        cosine_epochs=cosine_epochs,
        eta_min=eta_min,
        tail_lr=tail_lr,
        last_epoch=last_epoch,
    )
    return scheduler, cosine_epochs, tail_epochs, tail_lr

create_pytorch_cosine_lr_scheduler(optimizer, training, *, last_epoch=-1)

PyTorch CosineAnnealingLR; if num_epochs > T_max, LR restarts after each minimum.

Source code in oriented_det/train/utils.py

def create_pytorch_cosine_lr_scheduler(
    optimizer: "optim.Optimizer",
    training: Any,
    *,
    last_epoch: int = -1,
) -> Tuple[Any, int, float]:
    """PyTorch ``CosineAnnealingLR``; if ``num_epochs`` > ``T_max``, LR restarts after each minimum."""
    _require_torch()
    t_max = resolve_pytorch_cosine_t_max(training)
    eta_min = float(getattr(training, "lr_scheduler_cosine_eta_min", 1e-6))
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=t_max,
        eta_min=eta_min,
        last_epoch=last_epoch,
    )
    return scheduler, t_max, eta_min

effective_eval_metric_thresholds(config)

Score / per-class / IoU thresholds passed to validation and mAP.

• score_threshold: production.score_threshold when set, else evaluation.score_threshold.
• iou_threshold: always evaluation.iou_threshold (mAP / val matching).
• per_class_score_threshold: start from evaluation entries, then production updates keys.

Source code in oriented_det/train/config.py

def effective_eval_metric_thresholds(
    config: "TrainingExperimentConfig",
) -> Tuple[float, Optional[Dict[str, float]], float]:
    """Score / per-class / IoU thresholds passed to validation and mAP.

    - ``score_threshold``: ``production.score_threshold`` when set, else ``evaluation.score_threshold``.
    - ``iou_threshold``: always ``evaluation.iou_threshold`` (mAP / val matching).
    - ``per_class_score_threshold``: start from ``evaluation`` entries, then ``production`` updates keys.
    """
    ev = config.evaluation
    sc = float(ev.score_threshold)
    pc: Optional[Dict[str, float]] = None
    if ev.per_class_score_threshold:
        pc = {str(k): float(v) for k, v in ev.per_class_score_threshold.items()}
    iou = float(ev.iou_threshold)
    inf = getattr(config, "production", None)
    if inf is None:
        return sc, pc, iou
    if inf.score_threshold is not None:
        sc = float(inf.score_threshold)
    if inf.per_class_score_threshold is not None:
        merged: Dict[str, float] = dict(pc) if pc else {}
        merged.update({str(k): float(v) for k, v in inf.per_class_score_threshold.items()})
        pc = merged if merged else None
    return sc, pc, iou

evaluate(model, data_loader, device, *, metric_tracker=None, writer=None, epoch=None, class_map=None, class_names=None, score_threshold=0.05, per_class_score_threshold=None, vis_score_threshold=None, iou_threshold=0.5, extended_gt_metrics=False, compute_map=True, log_images=True, max_images_to_log=4, fixed_image_index=0, log_random_image=True, image_pool_max_size=32, log_debug_anchors_proposals=False, normalize_mean=None, normalize_std=None, vis_image_size=None, rank=None, progress_stream=None, debug=False, eval_use_exact_rotated_iou=True, compute_matching_metrics=True)

Evaluate model on validation set.

Note: When using DDP, this function bypasses DDP synchronization during forward pass by using model.module. This prevents deadlocks when only rank 0 runs validation while other ranks wait at barriers.

Parameters:

Name	Type	Description	Default
model	Module	Model to evaluate (may be DDP-wrapped)	required
data_loader	DataLoader	DataLoader for validation data	required
device	device	Device to evaluate on	required
metric_tracker	Optional[MetricTracker]	Optional metric tracker	None
writer	Optional[Any]	Optional TensorBoard SummaryWriter for logging	None
epoch	Optional[int]	Optional epoch number for logging	None
class_map	Optional[Dict[str, int]]	Optional mapping from class names to class IDs (for converting predictions)	None
class_names	Optional[Sequence[str]]	Optional list of class names (for evaluation)	None
score_threshold	float	Minimum confidence score for detections (metrics and matching)	0.05
vis_score_threshold	Optional[float]	Min score for boxes drawn in TensorBoard images; None = use score_threshold	None
iou_threshold	float	IoU threshold for mAP calculation	0.5
compute_map	bool	Whether to compute mAP (can be slow with many detections)	True
log_images	bool	Whether to log prediction images to TensorBoard	True
max_images_to_log	int	Maximum number of images to log per epoch (to avoid slowing down training)	4
fixed_image_index	int	Index in the validation order of the image to always log (0 = first image of first batch).	0
log_random_image	bool	If True, also log one randomly selected image from the pool (in addition to the fixed one).	True
image_pool_max_size	int	Maximum number of validation images to keep in the pool for random selection.	32
log_debug_anchors_proposals	bool	If True and model supports it, log one randomly selected validation image with anchors and one with RPN proposals to TensorBoard (for debugging). Random selection avoids always using the first image, which may have no ground truth.	False
normalize_mean	Optional[list[float]]	Normalization mean values [R, G, B] for denormalization. Defaults to ImageNet mean.	None
normalize_std	Optional[list[float]]	Normalization std values [R, G, B] for denormalization. Defaults to ImageNet std.	None
vis_image_size	Optional[Tuple[int, int]]	Optional (height, width) of the content region for TensorBoard images. When set, the image tensor is cropped to this size before drawing boxes so that annotations (in preprocessing target space) align with the displayed image. Default None uses the full tensor.	None
rank	Optional[int]	Optional rank for DDP (only rank 0 logs/prints, None for single-GPU)	None
progress_stream	Optional[Any]	If set (e.g. original stderr when stdout is teed to a log file), tqdm writes here so the progress bar is not duplicated in the log file.	None

Returns:

Type	Description
Dict[str, float]	Dictionary of evaluation metrics including mAP, per-class AP, GT classes, pred classes, and accuracy

Source code in oriented_det/train/engine.py

@torch.no_grad()
def evaluate(
    model: nn.Module,
    data_loader: DataLoader,
    device: torch.device,
    *,
    metric_tracker: Optional[MetricTracker] = None,
    writer: Optional[Any] = None,
    epoch: Optional[int] = None,
    class_map: Optional[Dict[str, int]] = None,
    class_names: Optional[Sequence[str]] = None,
    score_threshold: float = 0.05,
    per_class_score_threshold: Optional[Dict[str, float]] = None,
    vis_score_threshold: Optional[float] = None,
    iou_threshold: float = 0.5,
    extended_gt_metrics: bool = False,
    compute_map: bool = True,
    log_images: bool = True,
    max_images_to_log: int = 4,
    fixed_image_index: int = 0,
    log_random_image: bool = True,
    image_pool_max_size: int = 32,
    log_debug_anchors_proposals: bool = False,
    normalize_mean: Optional[list[float]] = None,
    normalize_std: Optional[list[float]] = None,
    vis_image_size: Optional[Tuple[int, int]] = None,
    rank: Optional[int] = None,
    progress_stream: Optional[Any] = None,
    debug: bool = False,
    eval_use_exact_rotated_iou: bool = True,
    compute_matching_metrics: bool = True,
) -> Dict[str, float]:
    """Evaluate model on validation set.

    Note: When using DDP, this function bypasses DDP synchronization during forward pass
    by using model.module. This prevents deadlocks when only rank 0 runs validation
    while other ranks wait at barriers.

    Args:
        model: Model to evaluate (may be DDP-wrapped)
        data_loader: DataLoader for validation data
        device: Device to evaluate on
        metric_tracker: Optional metric tracker
        writer: Optional TensorBoard SummaryWriter for logging
        epoch: Optional epoch number for logging
        class_map: Optional mapping from class names to class IDs (for converting predictions)
        class_names: Optional list of class names (for evaluation)
        score_threshold: Minimum confidence score for detections (metrics and matching)
        vis_score_threshold: Min score for boxes drawn in TensorBoard images; None = use score_threshold
        iou_threshold: IoU threshold for mAP calculation
        compute_map: Whether to compute mAP (can be slow with many detections)
        log_images: Whether to log prediction images to TensorBoard
        max_images_to_log: Maximum number of images to log per epoch (to avoid slowing down training)
        fixed_image_index: Index in the validation order of the image to always log (0 = first image of first batch).
        log_random_image: If True, also log one randomly selected image from the pool (in addition to the fixed one).
        image_pool_max_size: Maximum number of validation images to keep in the pool for random selection.
        log_debug_anchors_proposals: If True and model supports it, log one randomly selected validation image with anchors and one with RPN proposals to TensorBoard (for debugging). Random selection avoids always using the first image, which may have no ground truth.
        normalize_mean: Normalization mean values [R, G, B] for denormalization. Defaults to ImageNet mean.
        normalize_std: Normalization std values [R, G, B] for denormalization. Defaults to ImageNet std.
        vis_image_size: Optional (height, width) of the content region for TensorBoard images. When set, the image tensor is cropped to this size before drawing boxes so that annotations (in preprocessing target space) align with the displayed image. Default None uses the full tensor.
        rank: Optional rank for DDP (only rank 0 logs/prints, None for single-GPU)
        progress_stream: If set (e.g. original stderr when stdout is teed to a log file),
            tqdm writes here so the progress bar is not duplicated in the log file.

    Returns:
        Dictionary of evaluation metrics including mAP, per-class AP, GT classes, pred classes, and accuracy
    """
    _require_torch()
    import sys
    # Extract underlying model if DDP-wrapped (avoids DDP sync during validation)
    # CRITICAL: Set eval mode on underlying model, not DDP-wrapped model
    # Calling .eval() on DDP model can trigger synchronization operations
    # Also disable DDP's require_backward_grad_sync if DDP-wrapped to prevent any sync attempts
    eval_model = model.module if hasattr(model, 'module') else model
    eval_model.eval()

    # If DDP-wrapped, disable gradient synchronization during validation
    # This ensures no NCCL operations occur even if something tries to sync
    original_sync_state = None
    if hasattr(model, 'require_backward_grad_sync'):
        # Temporarily disable gradient sync (though we're in no_grad anyway)
        original_sync_state = model.require_backward_grad_sync
        model.require_backward_grad_sync = False

    metric_tracker = metric_tracker or MetricTracker()

    # Create reverse mapping from class_id to class_name
    id_to_class = {}
    if class_map is not None:
        id_to_class = {v: k for k, v in class_map.items()}

    # Collect all detections and ground truths
    all_detections: Dict[str, List[Detection]] = {}
    all_ground_truths: Dict[str, List[GroundTruth]] = {}

    # Track classes and counts
    gt_classes_set = set()
    pred_classes_set = set()
    total_correct = 0
    total_predictions = 0
    total_ground_truths = 0
    # Log-only diagnostics: how many GTs are covered before/after eval score filtering.
    gt_covered_pre_eval_threshold = 0
    gt_covered_post_eval_threshold = 0
    gt_lost_by_eval_threshold = 0
    max_detection_score = 0.0
    sum_detection_scores = 0.0
    num_detection_scores = 0

    # Pool of (vis_tensor, image_id) for TensorBoard: we pick fixed + random from this
    image_pool: Optional[list[tuple[Any, str]]] = [] if (writer is not None and log_images) else None
    # Debug images: anchors and proposals from first batch (only when log_debug_anchors_proposals and model supports it)
    debug_anchors_vis: Optional[torch.Tensor] = None
    debug_proposals_vis: Optional[torch.Tensor] = None

    # Create progress bar for validation (only on rank 0 or when rank is None).
    # When progress_stream is set, tqdm writes there so the log file doesn't get every update.
    num_batches = len(data_loader)
    # When logging debug anchors/proposals, pick a random batch so we don't always use the first image (which may have no GT).
    debug_batch_idx = random.randint(0, num_batches - 1) if (log_debug_anchors_proposals and num_batches > 0) else -1
    pbar = None
    if tqdm is not None and (rank is None or rank == 0):
        tqdm_file = progress_stream if progress_stream is not None else sys.stderr
        pbar = tqdm(
            total=num_batches,
            desc=f"Validation",
            unit="batch",
            ncols=120,
            leave=True,
            file=tqdm_file,
            disable=False,
        )

    for batch_idx, batch in enumerate(data_loader):
        step_start = time.time()
        try:
            # Unpack batch
            if isinstance(batch, (list, tuple)) and len(batch) == 2:
                images, targets = batch
            else:
                images, targets = batch.get("images", []), batch.get("targets", [])

            # Move to device
            if isinstance(images, list):
                images = [img.to(device) if torch.is_tensor(img) else img for img in images]
            elif torch.is_tensor(images):
                images = images.to(device)

            # Optional: ask model to return anchors and proposals for TensorBoard (random batch when log_debug_anchors_proposals)
            if batch_idx == debug_batch_idx and writer is not None and log_images and log_debug_anchors_proposals:
                setattr(eval_model, '_return_anchors_proposals', True)

            # Forward pass
            # CRITICAL: When all ranks participate in validation (DistributedSampler),
            # we can use the DDP-wrapped model directly. DDP will handle synchronization
            # correctly since all ranks are participating.
            # When only rank 0 validates, use eval_model (model.module) to bypass DDP.
            # Check if this is distributed validation (all ranks have data) or single-rank
            if rank is not None and dist is not None and dist.is_initialized():
                # Distributed validation: all ranks participate, use DDP model
                outputs = model(images)
            else:
                # Single-rank validation: use underlying model to avoid DDP sync
                outputs = eval_model(images)

            # Process each image in the batch
            if not isinstance(outputs, list):
                outputs = [outputs]
            if not isinstance(targets, list):
                targets = [targets]

            # Ensure images is a list for iteration
            if not isinstance(images, list):
                images = [images]
            # Pick random image index for debug anchors/proposals when this is the chosen batch
            debug_img_idx = -1
            if log_debug_anchors_proposals and batch_idx == debug_batch_idx and debug_anchors_vis is None:
                n_imgs = len(images)
                debug_img_idx = random.randint(0, n_imgs - 1) if n_imgs > 0 else 0

            # Track detections in current batch
            batch_detections = 0
            batch_detections_050 = 0  # count at score >= 0.5 (more informative when model is overconfident)

            for i, (output, target, image_tensor) in enumerate(zip(outputs, targets, images)):
                # Get image_id
                image_id = target.get("image_id", f"batch_{batch_idx}_img_{i}")
                if isinstance(image_id, torch.Tensor):
                    image_id = str(image_id.item())
                elif not isinstance(image_id, str):
                    image_id = str(image_id)
                display_name = target.get("image_filename")
                if display_name is None:
                    caption = image_id
                elif isinstance(display_name, torch.Tensor):
                    caption = str(display_name.item())
                elif not isinstance(display_name, str):
                    caption = str(display_name)
                else:
                    caption = display_name

                # Extract predictions (only rboxes format supported)
                rboxes = output.get("rboxes", [])
                scores = output.get("scores", torch.zeros(0))
                labels = output.get("labels", torch.zeros(0, dtype=torch.int64))

                # Ensure rboxes is a list of RBox objects
                if isinstance(rboxes, torch.Tensor) and len(rboxes) > 0:
                    # Convert tensor [N, 5] to list of RBox objects
                    if RBox is not None:
                        from ..models.utils import tensor_to_rboxes
                        rboxes = tensor_to_rboxes(rboxes)
                    else:
                        rboxes = []
                elif not isinstance(rboxes, list):
                    rboxes = []

                # Count detections at score >= 0.5 (for reporting; helps when eval threshold is low and model outputs 100)
                if len(scores) > 0:
                    if torch.is_tensor(scores):
                        batch_detections_050 += (scores >= 0.5).sum().item()
                        max_detection_score = max(max_detection_score, float(scores.max().item()))
                        sum_detection_scores += float(scores.sum().item())
                        num_detection_scores += scores.numel()
                    else:
                        batch_detections_050 += sum(1 for s in scores if float(s) >= 0.5)
                        for s in scores:
                            v = float(s)
                            max_detection_score = max(max_detection_score, v)
                            sum_detection_scores += v
                            num_detection_scores += 1
                # Keep a copy of raw predictions before eval score filtering for diagnostics.
                raw_scores = scores
                raw_labels = labels
                raw_rboxes = list(rboxes) if isinstance(rboxes, list) else []

                # Filter by score threshold (global and optional per-class)
                if len(scores) > 0:
                    if per_class_score_threshold and id_to_class:
                        mask = scores_labels_pass_threshold(
                            scores,
                            labels,
                            score_threshold,
                            per_class_score_threshold,
                            id_to_class,
                        )
                    else:
                        mask = scores >= score_threshold
                    scores = scores[mask]
                    labels = labels[mask]
                    if isinstance(rboxes, list) and len(rboxes) > 0:
                        rboxes = [rbox for rbox, keep in zip(rboxes, mask) if keep]

                # When padding was applied, detections are in padded image coords; GTs are in content coords. Clamp detection centers to content rect for correct IoU.
                def _clamp_rbox_to_content(rb, c_h: float, c_w: float):
                    cx_c = max(0.0, min(c_w, rb.cx))
                    cy_c = max(0.0, min(c_h, rb.cy))
                    return RBox(cx_c, cy_c, rb.width, rb.height, rb.angle)

                clamp_rbox = lambda r: r
                content_size = target.get("content_size")
                if content_size is not None and RBox is not None:
                    content_h, content_w = float(content_size[0]), float(content_size[1])
                    im_h = float(image_tensor.shape[-2])
                    im_w = float(image_tensor.shape[-1])
                    if (content_h, content_w) != (im_h, im_w):
                        clamp_rbox = lambda r, ch=content_h, cw=content_w: _clamp_rbox_to_content(r, ch, cw)

                # Convert predictions to Detection objects
                detections = []
                raw_detections = []
                if RBox is not None and Detection is not None:
                    # Build pre-filter detections (for log diagnostics only)
                    num_raw_dets = len(raw_scores)
                    if isinstance(raw_rboxes, list) and len(raw_rboxes) == num_raw_dets:
                        for rbox, score, label_id in zip(raw_rboxes, raw_scores, raw_labels):
                            label_id_int = int(label_id.item()) if torch.is_tensor(label_id) else int(label_id)
                            class_name = id_to_class.get(label_id_int, f"class_{label_id_int}")
                            raw_detections.append(Detection(
                                rbox=clamp_rbox(rbox),
                                score=float(score.item()) if torch.is_tensor(score) else float(score),
                                class_id=label_id_int,
                                class_name=class_name,
                                image_id=image_id,
                            ))

                    # Ensure all lists have the same length
                    num_dets = len(scores)
                    if isinstance(rboxes, list) and len(rboxes) == num_dets:
                        for j, (rbox, score, label_id) in enumerate(zip(rboxes, scores, labels)):
                            label_id_int = int(label_id.item()) if torch.is_tensor(label_id) else int(label_id)
                            class_name = id_to_class.get(label_id_int, f"class_{label_id_int}")
                            pred_classes_set.add(class_name)
                            detections.append(Detection(
                                rbox=clamp_rbox(rbox),
                                score=float(score.item()) if torch.is_tensor(score) else float(score),
                                class_id=label_id_int,
                                class_name=class_name,
                                image_id=image_id,
                            ))
                            total_predictions += 1

                batch_detections += len(detections)
                # Store raw (pre-score-threshold) detections so DDP gather can recompute pre/post stats on merged data
                all_detections[image_id] = raw_detections

                # Extract ground truths (only rboxes format supported)
                gt_labels = target.get("labels", torch.zeros(0, dtype=torch.int64))
                gt_rboxes = target.get("rboxes", None)
                gt_labels_ignore = target.get("labels_ignore", torch.zeros(0, dtype=torch.int64))
                gt_rboxes_ignore = target.get("rboxes_ignore", None)


                # Convert ground truths to GroundTruth objects
                ground_truths = []
                if gt_rboxes is not None:
                    # Handle rboxes (can be tensor [N, 5] or list of RBox objects)
                    if isinstance(gt_rboxes, torch.Tensor):
                        # Convert tensor [N, 5] to list of RBox objects
                        if RBox is not None:
                            from ..models.utils import tensor_to_rboxes
                            gt_rboxes_list = tensor_to_rboxes(gt_rboxes)
                        else:
                            gt_rboxes_list = []
                    elif isinstance(gt_rboxes, list):
                        gt_rboxes_list = gt_rboxes
                    else:
                        gt_rboxes_list = []

                    # Use rboxes if available
                    if len(gt_rboxes_list) > 0:
                        for rbox, label_id in zip(gt_rboxes_list, gt_labels):
                            if RBox is None or GroundTruth is None:
                                break
                            label_id_int = int(label_id.item()) if torch.is_tensor(label_id) else int(label_id)
                            class_name = id_to_class.get(label_id_int, f"class_{label_id_int}")
                            gt_classes_set.add(class_name)
                            ground_truths.append(GroundTruth(
                                rbox=rbox,
                                class_id=label_id_int,
                                class_name=class_name,
                                difficult=0,
                                image_id=image_id,
                            ))

                # Optional ignore GTs (e.g. DOTA difficult objects when difficult_strategy="ignore").
                if gt_rboxes_ignore is not None:
                    if isinstance(gt_rboxes_ignore, torch.Tensor):
                        if RBox is not None:
                            from ..models.utils import tensor_to_rboxes
                            gt_rboxes_ign_list = tensor_to_rboxes(gt_rboxes_ignore)
                        else:
                            gt_rboxes_ign_list = []
                    elif isinstance(gt_rboxes_ignore, list):
                        gt_rboxes_ign_list = gt_rboxes_ignore
                    else:
                        gt_rboxes_ign_list = []

                    if len(gt_rboxes_ign_list) > 0:
                        for rbox, label_id in zip(gt_rboxes_ign_list, gt_labels_ignore):
                            if RBox is None or GroundTruth is None:
                                break
                            label_id_int = int(label_id.item()) if torch.is_tensor(label_id) else int(label_id)
                            class_name = id_to_class.get(label_id_int, f"class_{label_id_int}")
                            gt_classes_set.add(class_name)
                            ground_truths.append(GroundTruth(
                                rbox=rbox,
                                class_id=label_id_int,
                                class_name=class_name,
                                difficult=1,
                                image_id=image_id,
                            ))
                # Note: No fallback for "boxes" format - only "rboxes" is supported

                all_ground_truths[image_id] = ground_truths
                total_ground_truths += len(ground_truths)

                # Add visualization to pool for TensorBoard (capped to avoid OOM).
                # When log_debug_anchors_proposals is True, only add the same random image we use for debug (anchors/proposals) so predictions show that one image.
                add_to_pool = image_pool is not None and len(image_pool) < image_pool_max_size
                if add_to_pool and log_debug_anchors_proposals:
                    add_to_pool = batch_idx == debug_batch_idx and i == debug_img_idx
                if add_to_pool:
                    try:
                        # Use content region so image and boxes share the same coordinate system.
                        # (Boxes are in preprocessing target size; tensor may be padded.)
                        vis_tensor = image_tensor
                        if vis_image_size is not None:
                            vh, vw = vis_image_size[0], vis_image_size[1]
                            _, th, tw = image_tensor.shape
                            if th >= vh and tw >= vw:
                                vis_tensor = image_tensor[:, :vh, :vw].clone()
                        # Convert tensor to PIL Image
                        pil_image = _tensor_to_pil_image(
                            vis_tensor,
                            mean=normalize_mean,
                            std=normalize_std,
                        )
                        # Boxes to draw: use vis_score_threshold if set, else score_threshold (detections)
                        vis_thresh = vis_score_threshold if vis_score_threshold is not None else score_threshold
                        detections_for_vis = [d for d in raw_detections if d.score >= vis_thresh]
                        # Create visualization
                        vis_image = _visualize_predictions(
                            pil_image.copy(),
                            detections_for_vis,
                            ground_truths,
                            class_names=class_names,
                            caption=caption,
                        )
                        # Convert back to tensor for TensorBoard (C, H, W) format
                        if np is not None:
                            vis_array = np.array(vis_image).transpose(2, 0, 1)  # (H, W, C) -> (C, H, W)
                            vis_tensor = torch.from_numpy(vis_array).float() / 255.0
                            image_pool.append((vis_tensor, image_id))
                        # Debug: anchors and proposals from a random image (model may set output["anchors"], output["proposals"])
                        if batch_idx == debug_batch_idx and i == debug_img_idx and log_debug_anchors_proposals and debug_anchors_vis is None:
                            anchors_t = output.get("anchors")
                            proposals_t = output.get("proposals")
                            if anchors_t is not None and isinstance(anchors_t, torch.Tensor) and anchors_t.numel() > 0:
                                try:
                                    pil_anchors = _tensor_to_pil_image(image_tensor, mean=normalize_mean, std=normalize_std)
                                    vis_anchors = _visualize_boxes(pil_anchors.copy(), anchors_t, color=(255, 200, 0), max_boxes=500)
                                    vis_anchors = _draw_filename_caption(vis_anchors, caption)
                                    if np is not None:
                                        arr = np.array(vis_anchors).transpose(2, 0, 1)
                                        debug_anchors_vis = torch.from_numpy(arr).float() / 255.0
                                except Exception:
                                    pass
                            if proposals_t is not None and isinstance(proposals_t, torch.Tensor) and proposals_t.numel() > 0:
                                try:
                                    pil_proposals = _tensor_to_pil_image(image_tensor, mean=normalize_mean, std=normalize_std)
                                    vis_proposals = _visualize_boxes(pil_proposals.copy(), proposals_t, color=(0, 200, 255), max_boxes=500)
                                    vis_proposals = _draw_filename_caption(vis_proposals, caption)
                                    if np is not None:
                                        arr = np.array(vis_proposals).transpose(2, 0, 1)
                                        debug_proposals_vis = torch.from_numpy(arr).float() / 255.0
                                except Exception:
                                    pass
                    except Exception as e:
                        # Silently skip if visualization fails (e.g., missing dependencies)
                        pass

            # Track basic metrics (detections per image, not total per batch)
            num_images_in_batch = len(outputs) if isinstance(outputs, list) else 1
            dets_per_image = batch_detections / num_images_in_batch if num_images_in_batch > 0 else 0.0
            dets_per_image_050 = batch_detections_050 / num_images_in_batch if num_images_in_batch > 0 else 0.0
            metrics = {"num_detections": dets_per_image, "num_detections_050": dets_per_image_050}
            metric_tracker.update(metrics)

            # Clear debug flag after the batch we used so subsequent batches don't return anchors/proposals
            if batch_idx == debug_batch_idx and getattr(eval_model, '_return_anchors_proposals', False):
                setattr(eval_model, '_return_anchors_proposals', False)

            # Update progress bar
            elapsed = time.time() - step_start
            metric_tracker.update_time(elapsed)

            if pbar is not None:
                # Update progress bar every batch
                pbar.update(1)

                # Update metrics display periodically (every 10 batches or at the end)
                if (batch_idx + 1) % 10 == 0 or (batch_idx + 1) == num_batches:
                    summary = metric_tracker.get_summary(window=10)
                    # Average detections per batch over the window
                    avg_dets = summary.get('num_detections', 0)
                    avg_time = summary.get('time_per_step', elapsed)

                    # Calculate speed and ETA
                    batches_per_sec = 1.0 / avg_time if avg_time > 0 else 0.0
                    remaining_batches = num_batches - (batch_idx + 1)
                    eta_seconds = remaining_batches * avg_time if avg_time > 0 else 0.0

                    # Format ETA
                    if eta_seconds >= 3600:
                        eta_str = f"{int(eta_seconds // 3600):d}:{int((eta_seconds % 3600) // 60):02d}:{int(eta_seconds % 60):02d}"
                    else:
                        eta_str = f"{int(eta_seconds // 60):d}:{int(eta_seconds % 60):02d}"

                    num_dets_str = f"Dets/img: {avg_dets:.1f}"
                    speed_str = f"{batches_per_sec:.2f}it/s"
                    eta_str_tqdm = f"ETA: {eta_str}"
                    pbar.set_postfix_str(f"{num_dets_str} | {speed_str} | {eta_str_tqdm}")
            elif rank is None or rank == 0:
                # Fallback: print progress with speed and ETA if tqdm is not available
                if (batch_idx + 1) % 10 == 0 or (batch_idx + 1) == num_batches:
                    summary = metric_tracker.get_summary(window=10)
                    avg_dets = summary.get('num_detections', 0)
                    avg_time = summary.get('time_per_step', elapsed)

                    # Calculate speed and ETA
                    batches_per_sec = 1.0 / avg_time if avg_time > 0 else 0.0
                    remaining_batches = num_batches - (batch_idx + 1)
                    eta_seconds = remaining_batches * avg_time if avg_time > 0 else 0.0

                    # Format ETA
                    if eta_seconds >= 3600:
                        eta_str = f"{int(eta_seconds // 3600):d}:{int((eta_seconds % 3600) // 60):02d}:{int(eta_seconds % 60):02d}"
                    else:
                        eta_str = f"{int(eta_seconds // 60):d}:{int(eta_seconds % 60):02d}"

                    print(
                        f"Validation [{batch_idx + 1}/{num_batches}] "
                        f"Dets/img: {avg_dets:.1f} | "
                        f"Speed: {batches_per_sec:.2f}it/s | "
                        f"ETA: {eta_str}"
                    )
                    import sys
                    sys.stdout.flush()

        except Exception as e:
            if pbar is not None:
                pbar.close()
            if rank is None or rank == 0:
                print(f"Error in evaluation batch {batch_idx}: {e}")
                import traceback
                traceback.print_exc()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            continue

    # Close progress bar
    if pbar is not None:
        pbar.close()

    # Single line to stdout (so log file gets one line when stdout is teed)
    if (rank is None or rank == 0) and num_batches > 0:
        print(f"Validation complete: {num_batches} batches", flush=True)

    # DDP: gather all_detections and all_ground_truths from all ranks to rank 0, merge, and recompute stats on full val set
    did_gather = False
    val_stats_extra: Dict[str, Any] = {}
    if dist is not None and dist.is_initialized() and rank is not None:
        world_size = dist.get_world_size()
        if world_size > 1:
            payload = (list(all_detections.items()), list(all_ground_truths.items()))
            object_list = [None] * world_size if rank == 0 else None
            # Finish GPU work before collectives; helps surface async CUDA errors here, not in gather.
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            # Ensure every rank reaches gather together after validation.
            dist.barrier()
            gather_group = _get_gloo_gather_object_group()
            gather_success = False
            try:
                if gather_group is not None:
                    try:
                        dist.gather_object(payload, object_list, dst=0, group=gather_group)
                    except TypeError:
                        # PyTorch without `group=` on gather_object — use default process group
                        dist.gather_object(payload, object_list, dst=0)
                else:
                    dist.gather_object(payload, object_list, dst=0)
                gather_success = True
            except AttributeError:
                # Older PyTorch may not have gather_object; keep shard-only metrics
                object_list = None
            except RuntimeError as e:
                # Last resort if Gloo gather still fails: log and fall back to per-rank shard metrics.
                if rank == 0:
                    print(
                        f"Warning: dist.gather_object failed ({e!r}); "
                        "using per-rank validation shard only (mAP/stats may be incomplete).",
                        flush=True,
                    )
                object_list = None
            if gather_success:
                if rank == 0:
                    did_gather = True
                    merged_detections = {}
                    merged_ground_truths = {}
                    for (det_items, gt_items) in object_list:
                        for k, v in dict(det_items).items():
                            merged_detections[k] = merged_detections.get(k, []) + list(v)
                        for k, v in dict(gt_items).items():
                            merged_ground_truths[k] = merged_ground_truths.get(k, []) + list(v)
                    all_detections = merged_detections
                    all_ground_truths = merged_ground_truths
                    stats = _compute_val_stats_from_dicts(
                        all_detections,
                        all_ground_truths,
                        score_threshold,
                        iou_threshold,
                        id_to_class,
                        per_class_score_threshold,
                        extended_gt_metrics,
                        compute_matching_metrics=compute_matching_metrics,
                        use_exact_rotated_iou=eval_use_exact_rotated_iou,
                        device=device,
                    )
                    total_ground_truths = stats["total_ground_truths"]
                    gt_covered_pre_eval_threshold = stats["gt_covered_pre_eval_threshold"]
                    gt_covered_post_eval_threshold = stats["gt_covered_post_eval_threshold"]
                    gt_lost_by_eval_threshold = stats["gt_lost_by_eval_threshold"]
                    total_correct = stats["total_correct"]
                    total_predictions = stats["total_predictions"]
                    gt_classes_set = stats["gt_classes_set"]
                    pred_classes_set = stats["pred_classes_set"]
                    max_detection_score = stats["max_detection_score"]
                    sum_detection_scores = stats["sum_detection_scores"]
                    num_detection_scores = stats["num_detection_scores"]
                    val_stats_extra = stats
                else:
                    did_gather = True
                    total_ground_truths = 0
                    gt_covered_pre_eval_threshold = 0
                    gt_covered_post_eval_threshold = 0
                    gt_lost_by_eval_threshold = 0
                    total_correct = 0
                    total_predictions = 0
                    gt_classes_set = set()
                    pred_classes_set = set()
                    max_detection_score = 0.0
                    sum_detection_scores = 0.0
                    num_detection_scores = 0

    # Single-process / single-GPU: recompute stats from full dicts (includes GT–IoU diagnostics).
    if not did_gather and (rank is None or rank == 0):
        val_stats_extra = _compute_val_stats_from_dicts(
            all_detections,
            all_ground_truths,
            score_threshold,
            iou_threshold,
            id_to_class,
            per_class_score_threshold,
            extended_gt_metrics,
            compute_matching_metrics=compute_matching_metrics,
            use_exact_rotated_iou=eval_use_exact_rotated_iou,
            device=device,
        )
        total_ground_truths = val_stats_extra["total_ground_truths"]
        gt_covered_pre_eval_threshold = val_stats_extra["gt_covered_pre_eval_threshold"]
        gt_covered_post_eval_threshold = val_stats_extra["gt_covered_post_eval_threshold"]
        gt_lost_by_eval_threshold = val_stats_extra["gt_lost_by_eval_threshold"]
        total_correct = val_stats_extra["total_correct"]
        total_predictions = val_stats_extra["total_predictions"]
        gt_classes_set = val_stats_extra["gt_classes_set"]
        pred_classes_set = val_stats_extra["pred_classes_set"]
        max_detection_score = val_stats_extra["max_detection_score"]
        sum_detection_scores = val_stats_extra["sum_detection_scores"]
        num_detection_scores = val_stats_extra["num_detection_scores"]

    # Compute DOTA metrics
    summary = metric_tracker.get_summary()

    # Compute mAP if evaluation utilities are available and requested (only on rank 0 when we gathered)
    if compute_map and compute_oriented_map is not None and Detection is not None and GroundTruth is not None:
        if did_gather and rank is not None and rank != 0:
            summary["mAP"] = -1.0
        else:
            try:
                # mAP uses score-filtered detections (all_detections stores raw for DDP stats)
                detections_for_map = {
                    k: filter_detections_by_score_threshold(
                        v, score_threshold, per_class_score_threshold, id_to_class
                    )
                    for k, v in all_detections.items()
                }
                total_dets = sum(len(dets) for dets in detections_for_map.values())
                total_gts = sum(len(gts) for gts in all_ground_truths.values())
                num_images = len(all_detections)
                if num_images == 0:
                    num_images = 1
                if rank is None or rank == 0:
                    print(f"\nComputing mAP (this may take a while)...")
                    print(
                        f"  Eval filters: score≥{score_threshold:.4f}, IoU≥{iou_threshold:.2f}"
                        + (
                            f", per-class floors on {len(per_class_score_threshold)} class(es)"
                            if per_class_score_threshold
                            else ""
                        )
                    )
                    print(f"  Images: {num_images:,}")
                    print(
                        f"  Total detections (post score filter): {total_dets:,} "
                        f"({total_dets/num_images:.1f} per image)"
                    )
                    print(f"  Total ground truths: {total_gts:,} ({total_gts/num_images:.1f} per image)")
                    if total_dets > 50000:
                        print(f"  Warning: Large number of detections may slow down mAP computation")
                        print(
                            f"  Consider increasing evaluation.score_threshold "
                            f"(current: {score_threshold}) for faster mAP matching"
                        )
                _map_t0 = time.perf_counter()
                mean_ap, class_aps, _class_metrics = compute_oriented_map(
                    detections=detections_for_map,
                    ground_truths=all_ground_truths,
                    iou_threshold=iou_threshold,
                    class_names=class_names,
                    device=device,
                    progress_stream=progress_stream,
                    use_exact_rotated_iou=eval_use_exact_rotated_iou,
                )
                _map_elapsed = time.perf_counter() - _map_t0
                summary["mAP"] = mean_ap
                for class_name, ap in class_aps.items():
                    summary[f"AP_{class_name}"] = ap
                if rank is None or rank == 0:
                    print(
                        f"  mAP computation completed in {_format_duration_hms(_map_elapsed)}. "
                        f"mAP: {mean_ap:.4f}"
                    )
                    if debug and class_aps:
                        print("  [debug] Per-class AP:", flush=True)
                        for cls_name, ap_val in sorted(class_aps.items(), key=lambda x: -x[1]):
                            print(f"    {cls_name}: {ap_val:.4f}", flush=True)
                        from collections import Counter
                        det_by_cls = Counter()
                        gt_by_cls = Counter()
                        for dets in detections_for_map.values():
                            for d in dets:
                                det_by_cls[d.class_name] += 1
                        for gts in all_ground_truths.values():
                            for g in gts:
                                gt_by_cls[g.class_name] += 1
                        print("  [debug] Detections per class:", dict(det_by_cls.most_common()), flush=True)
                        print("  [debug] Ground truths per class:", dict(gt_by_cls.most_common()), flush=True)
            except Exception as e:
                if rank is None or rank == 0:
                    print(f"Warning: Could not compute mAP: {e}")
                    import traceback
                    traceback.print_exc()
                summary["mAP"] = 0.0
    elif not compute_map:
        # mAP computation skipped
        summary["mAP"] = -1.0  # Use -1 to indicate it was skipped

    # Add class information
    summary["gt_classes"] = sorted(list(gt_classes_set))
    summary["pred_classes"] = sorted(list(pred_classes_set))

    # Compute accuracy (only when matching metrics were computed this epoch; otherwise the
    # counters are untouched zeros and printing them would look like a broken model).
    if compute_matching_metrics:
        if total_predictions > 0:
            summary["accuracy"] = total_correct / total_predictions
        else:
            summary["accuracy"] = 0.0
        summary["total_correct"] = total_correct

    summary["total_predictions"] = total_predictions
    summary["total_ground_truths"] = total_ground_truths
    summary["max_detection_score"] = max_detection_score
    if num_detection_scores > 0:
        summary["mean_detection_score"] = sum_detection_scores / num_detection_scores
    else:
        summary["mean_detection_score"] = 0.0

    # Log-only diagnostics (do not send to TensorBoard).
    summary["eval_score_threshold"] = score_threshold
    summary["eval_iou_threshold"] = iou_threshold
    if per_class_score_threshold:
        summary["eval_per_class_score_threshold"] = dict(per_class_score_threshold)
    if not compute_matching_metrics:
        pass  # GT-cover diagnostics skipped this epoch (same schedule as mAP)
    elif val_stats_extra:
        summary["log_only_gt_covered_pre_eval_threshold"] = val_stats_extra["gt_covered_pre_eval_threshold"]
        summary["log_only_gt_covered_post_eval_threshold"] = val_stats_extra["gt_covered_post_eval_threshold"]
        summary["log_only_gt_lost_by_eval_threshold"] = val_stats_extra["gt_lost_by_eval_threshold"]
        for _k, _v in val_stats_extra.items():
            if _k.startswith("log_only_gt_"):
                summary[_k] = _v
        if total_ground_truths > 0:
            summary["log_only_gt_cover_rate_pre_eval_threshold"] = (
                val_stats_extra["gt_covered_pre_eval_threshold"] / total_ground_truths
            )
            summary["log_only_gt_cover_rate_post_eval_threshold"] = (
                val_stats_extra["gt_covered_post_eval_threshold"] / total_ground_truths
            )
        else:
            summary["log_only_gt_cover_rate_pre_eval_threshold"] = 0.0
            summary["log_only_gt_cover_rate_post_eval_threshold"] = 0.0
    else:
        summary["log_only_gt_covered_pre_eval_threshold"] = gt_covered_pre_eval_threshold
        summary["log_only_gt_covered_post_eval_threshold"] = gt_covered_post_eval_threshold
        summary["log_only_gt_lost_by_eval_threshold"] = gt_lost_by_eval_threshold
        if total_ground_truths > 0:
            summary["log_only_gt_cover_rate_pre_eval_threshold"] = gt_covered_pre_eval_threshold / total_ground_truths
            summary["log_only_gt_cover_rate_post_eval_threshold"] = gt_covered_post_eval_threshold / total_ground_truths
        else:
            summary["log_only_gt_cover_rate_pre_eval_threshold"] = 0.0
            summary["log_only_gt_cover_rate_post_eval_threshold"] = 0.0

    # Debug: print GT cover rates and score stats to diagnose low mAP vs MMRotate.
    if debug and (rank is None or rank == 0):
        if total_ground_truths > 0:
            print(
                f"  [debug] GT cover (pre score thresh): {gt_covered_pre_eval_threshold}/{total_ground_truths} "
                f"({summary['log_only_gt_cover_rate_pre_eval_threshold']:.2%})",
                flush=True,
            )
            print(
                f"  [debug] GT cover (post score thresh): {gt_covered_post_eval_threshold}/{total_ground_truths} "
                f"({summary['log_only_gt_cover_rate_post_eval_threshold']:.2%}), "
                f"GT lost by threshold: {gt_lost_by_eval_threshold}",
                flush=True,
            )
        print(
            f"  [debug] Detection scores: mean={summary.get('mean_detection_score', 0):.4f}, "
            f"max={summary.get('max_detection_score', 0):.4f}, "
            f"total_predictions={total_predictions}, total_correct={total_correct}",
            flush=True,
        )

    # Build final list for TensorBoard: fixed image (selected) + optional random image(s)
    images_to_log: Optional[list[tuple[Any, str]]] = None
    if image_pool is not None and len(image_pool) > 0:
        images_to_log = []
        used_indices: set[int] = set()
        # Always include the fixed (selected) image: first image of first batch by default
        idx = min(fixed_image_index, len(image_pool) - 1)
        images_to_log.append(image_pool[idx])
        used_indices.add(idx)
        # Add one randomly selected image (different from fixed if pool has > 1)
        if log_random_image and len(image_pool) > 1:
            other_indices = [i for i in range(len(image_pool)) if i not in used_indices]
            if other_indices:
                random_idx = random.choice(other_indices)
                images_to_log.append(image_pool[random_idx])
                used_indices.add(random_idx)
        # Fill up to max_images_to_log with more random picks (without replacement)
        remaining = max_images_to_log - len(images_to_log)
        if remaining > 0:
            candidates = [i for i in range(len(image_pool)) if i not in used_indices]
            for _ in range(min(remaining, len(candidates))):
                pick = random.choice(candidates)
                candidates.remove(pick)
                images_to_log.append(image_pool[pick])

    # Log validation metrics to TensorBoard
    if writer is not None and epoch is not None:
        for key, value in summary.items():
            # Keep log-only diagnostics out of TensorBoard by design.
            if key.startswith("log_only_"):
                continue
            # Do not log mAP when it was not computed (sentinel -1)
            if key == "mAP" and value == -1.0:
                continue
            # Skip non-numeric values for scalar logging
            if isinstance(value, (int, float)):
                writer.add_scalar(f"val/{key}", value, epoch)
            elif isinstance(value, list):
                # Log class lists as text or count
                writer.add_scalar(f"val/{key}_count", len(value), epoch)

        # Log prediction images to TensorBoard
        if images_to_log is not None and len(images_to_log) > 0:
            try:
                # Create a grid of images for TensorBoard
                from torchvision.utils import make_grid
                vis_tensors = [img_tensor for img_tensor, _ in images_to_log]
                grid = make_grid(vis_tensors, nrow=min(2, len(vis_tensors)), padding=2)
                writer.add_image("val/predictions", grid, epoch)
            except Exception as e:
                # Silently skip if image logging fails (e.g., missing torchvision)
                pass
        # Log debug images: anchors and proposals (same random image as val/predictions when log_debug_anchors_proposals)
        if log_debug_anchors_proposals:
            try:
                if debug_anchors_vis is not None:
                    writer.add_image("val/debug_anchors", debug_anchors_vis, epoch)
                if debug_proposals_vis is not None:
                    writer.add_image("val/debug_proposals", debug_proposals_vis, epoch)
            except Exception:
                pass

    # Restore DDP sync state if we modified it
    if hasattr(model, 'require_backward_grad_sync') and original_sync_state is not None:
        model.require_backward_grad_sync = original_sync_state

    # DDP: rank 0 may run merge + stats + mAP for a long time after the val forward loop;
    # non-zero ranks used to return immediately and block on train()'s post-val barrier,
    # tripping Gloo's default 30-minute barrier timeout. Sync here so every rank leaves
    # evaluate() together (after rank 0's mAP / logging).
    if dist is not None and dist.is_initialized() and rank is not None:
        if dist.get_world_size() > 1:
            try:
                _gg = _get_gloo_gather_object_group()
                if _gg is not None:
                    dist.barrier(group=_gg)
                else:
                    dist.barrier()
            except TypeError:
                dist.barrier()

    return summary

format_cosine_with_tail_scheduler_description(cosine_epochs, tail_epochs, eta_min, tail_lr, warmup_steps)

Human-readable line for cosine_annealing_with_tail.

Source code in oriented_det/train/utils.py

def format_cosine_with_tail_scheduler_description(
    cosine_epochs: int,
    tail_epochs: int,
    eta_min: float,
    tail_lr: float,
    warmup_steps: int,
) -> str:
    """Human-readable line for ``cosine_annealing_with_tail``."""
    core = (
        f"CosineAnnealingWithFixedTailLR (cosine_epochs={cosine_epochs}, "
        f"tail_epochs={tail_epochs}, eta_min={eta_min:g}, tail_lr={tail_lr:g})"
    )
    if warmup_steps > 0:
        return f"{core} + warmup {warmup_steps} steps"
    return core

format_pytorch_cosine_scheduler_description(t_max, eta_min, warmup_steps, num_epochs)

Human-readable line for cosine_annealing (PyTorch scheduler).

Source code in oriented_det/train/utils.py

def format_pytorch_cosine_scheduler_description(
    t_max: int,
    eta_min: float,
    warmup_steps: int,
    num_epochs: int,
) -> str:
    """Human-readable line for ``cosine_annealing`` (PyTorch scheduler)."""
    restart_note = ""
    if num_epochs > t_max:
        restart_note = f"; SGDR-style restarts when num_epochs ({num_epochs}) > T_max"
    core = f"CosineAnnealingLR (T_max={t_max}, eta_min={eta_min:g}{restart_note})"
    if warmup_steps > 0:
        return f"{core} + warmup {warmup_steps} steps"
    return core

get_best_checkpoint_path(checkpoint_dir)

Return the path to the best checkpoint in the directory, if any.

Looks for checkpoint_best.pth (legacy) first, then best_*.pth (e.g. best_mAP_0.42.pth).

Source code in oriented_det/train/utils.py

def get_best_checkpoint_path(checkpoint_dir: Path) -> Optional[Path]:
    """Return the path to the best checkpoint in the directory, if any.

    Looks for checkpoint_best.pth (legacy) first, then best_*.pth (e.g. best_mAP_0.42.pth).
    """
    checkpoint_dir = Path(checkpoint_dir)
    legacy = checkpoint_dir / "checkpoint_best.pth"
    if legacy.exists():
        return legacy
    best_files = sorted(checkpoint_dir.glob("best_*.pth"))
    return best_files[0] if best_files else None

grouped_ce_alpha_for_epoch(epoch, *, enabled, schedule_type, start_epoch, end_epoch, power=1.0)

Return grouped-loss mix factor in [0, 1] (1 = fully grouped, 0 = fine CE only).

Source code in oriented_det/train/grouped_ce.py

def grouped_ce_alpha_for_epoch(
    epoch: int,
    *,
    enabled: bool,
    schedule_type: Optional[str],
    start_epoch: int,
    end_epoch: int,
    power: float = 1.0,
) -> float:
    """Return grouped-loss mix factor in [0, 1] (1 = fully grouped, 0 = fine CE only)."""
    if not enabled:
        return 0.0
    sched = (schedule_type or "step").strip().lower()
    if sched in ("step", "steps", "hard"):
        return 1.0 if epoch < end_epoch else 0.0
    if sched in ("linear_ramp", "ramp", "linear"):
        if end_epoch <= start_epoch:
            return 0.0
        if epoch <= start_epoch:
            return 1.0
        if epoch >= end_epoch:
            return 0.0
        t = (epoch - start_epoch) / float(end_epoch - start_epoch)
        return float((1.0 - t) ** power)
    return 0.0

profile_training_step(model, images, targets, optimizer, device, use_amp=False, profiler=None)

Profile a single training step.

Parameters:

Name	Type	Description	Default
model	Any	Model to train	required
images	Any	Input images	required
targets	Any	Target annotations	required
optimizer	Any	Optimizer	required
device	device	Device to train on	required
use_amp	bool	Use automatic mixed precision	False
profiler	Optional[TrainingProfiler]	Optional profiler context manager	None

Returns:

Type	Description
dict[str, float]	Dictionary with timing information

Source code in oriented_det/train/profiler.py

def profile_training_step(
    model: Any,
    images: Any,
    targets: Any,
    optimizer: Any,
    device: torch.device,
    use_amp: bool = False,
    profiler: Optional[TrainingProfiler] = None,
) -> dict[str, float]:
    """Profile a single training step.

    Args:
        model: Model to train
        images: Input images
        targets: Target annotations
        optimizer: Optimizer
        device: Device to train on
        use_amp: Use automatic mixed precision
        profiler: Optional profiler context manager

    Returns:
        Dictionary with timing information
    """
    _require_torch()

    import time

    timings = {}

    # Data loading time (if applicable)
    data_start = time.time()
    if isinstance(images, list):
        images = [img.to(device) if torch.is_tensor(img) else img for img in images]
    elif torch.is_tensor(images):
        images = images.to(device)
    timings["data_to_device"] = time.time() - data_start

    # Forward pass
    forward_start = time.time()
    if profiler is not None:
        with record_function("forward_pass"):
            if use_amp:
                try:
                    from torch.amp import autocast
                    with autocast('cuda'):
                        loss_dict = model(images, targets)
                except (ImportError, TypeError):
                    from torch.cuda.amp import autocast
                    with autocast():
                        loss_dict = model(images, targets)
            else:
                loss_dict = model(images, targets)
    else:
        if use_amp:
            try:
                from torch.amp import autocast
                with autocast('cuda'):
                    loss_dict = model(images, targets)
            except (ImportError, TypeError):
                from torch.cuda.amp import autocast
                with autocast():
                    loss_dict = model(images, targets)
        else:
            loss_dict = model(images, targets)
    timings["forward"] = time.time() - forward_start

    # Loss computation
    loss_start = time.time()
    total_loss = sum(v for v in loss_dict.values() if torch.is_tensor(v))
    timings["loss_computation"] = time.time() - loss_start

    # Backward pass
    backward_start = time.time()
    if profiler is not None:
        with record_function("backward_pass"):
            total_loss.backward()
    else:
        total_loss.backward()
    timings["backward"] = time.time() - backward_start

    # Optimizer step
    optimizer_start = time.time()
    if profiler is not None:
        with record_function("optimizer_step"):
            optimizer.step()
            optimizer.zero_grad()
    else:
        optimizer.step()
        optimizer.zero_grad()
    timings["optimizer"] = time.time() - optimizer_start

    timings["total"] = sum(timings.values())

    return timings

resolve_cosine_with_tail_lengths(training)

Cosine phase + fixed-LR tail lengths for cosine_annealing_with_tail (must sum to num_epochs).

Source code in oriented_det/train/utils.py

def resolve_cosine_with_tail_lengths(training: Any) -> Tuple[int, int]:
    """Cosine phase + fixed-LR tail lengths for ``cosine_annealing_with_tail`` (must sum to ``num_epochs``)."""
    num_epochs = int(training.num_epochs)
    cosine_epochs = _resolve_cosine_phase_epochs(training)
    tail_epochs_cfg = int(getattr(training, "lr_scheduler_cosine_tail_epochs", 0) or 0)
    if tail_epochs_cfg > 0:
        tail_epochs = tail_epochs_cfg
    else:
        tail_epochs = num_epochs - cosine_epochs
    if tail_epochs < 1:
        raise ValueError(
            f"cosine_annealing_with_tail needs tail_epochs >= 1; got cosine_epochs={cosine_epochs}, "
            f"num_epochs={num_epochs}. Set lr_scheduler_cosine_tail_epochs or shorten the cosine phase."
        )
    if cosine_epochs + tail_epochs != num_epochs:
        raise ValueError(
            f"cosine_epochs ({cosine_epochs}) + tail_epochs ({tail_epochs}) must equal "
            f"num_epochs ({num_epochs}) for cosine_annealing_with_tail"
        )
    return cosine_epochs, tail_epochs

resolve_inference_sliding_window_overlap_pixels(config)

Sliding-window overlap in pixels per axis for production inference (ratio-free).

Returns:

Type	Description
int	production.overlap_pixels when set (clamped to >= 0), else 200 (same default as
int	func:oriented_det.runtime.inference.run_inference_auto).

Source code in oriented_det/train/config.py

def resolve_inference_sliding_window_overlap_pixels(config: "TrainingExperimentConfig") -> int:
    """Sliding-window overlap in pixels per axis for production inference (ratio-free).

    Returns:
        ``production.overlap_pixels`` when set (clamped to ``>= 0``), else **200** (same default as
        :func:`oriented_det.runtime.inference.run_inference_auto`).
    """
    inf = getattr(config, "production", None)
    if inf is not None and getattr(inf, "overlap_pixels", None) is not None:
        return max(0, int(inf.overlap_pixels))
    return 200

resolve_pytorch_cosine_t_max(training)

T_max for :class:torch.optim.lr_scheduler.CosineAnnealingLR (PyTorch default, including restarts).

Source code in oriented_det/train/utils.py

def resolve_pytorch_cosine_t_max(training: Any) -> int:
    """``T_max`` for :class:`torch.optim.lr_scheduler.CosineAnnealingLR` (PyTorch default, including restarts)."""
    return _resolve_cosine_phase_epochs(training)

save_training_config(config, path)

Save training configuration to JSON file.

Source code in oriented_det/train/utils.py

def save_training_config(config: Dict[str, Any], path: str | Path) -> None:
    """Save training configuration to JSON file."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    with path.open("w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, default=str)

train(model, train_loader, optimizer, device, *, num_epochs=10, val_loader=None, lr_scheduler=None, checkpoint_manager=None, print_freq=10, gradient_accumulation_steps=1, max_grad_norm=None, use_amp=False, loss_weights=None, roi_class_weights=None, start_epoch=0, writer=None, class_map=None, class_names=None, eval_score_threshold=0.05, eval_per_class_score_threshold=None, eval_vis_score_threshold=None, eval_iou_threshold=0.5, eval_extended_gt_metrics=False, eval_compute_map_final=True, eval_compute_map_every_n_epochs=0, log_images=True, max_images_to_log=4, fixed_image_index=0, log_random_image=True, image_pool_max_size=128, log_debug_anchors_proposals=False, profiler=None, normalize_mean=None, normalize_std=None, vis_image_size=None, train_sampler=None, val_sampler=None, rank=None, progress_stream=None, debug=False, lr_scheduler_plateau_metric='total_loss', early_stop_patience=None, early_stop_metric='mAP', early_stop_min_delta=0.0, early_stop_higher_is_better=None, freeze_backbone_epochs=0, freeze_rpn_epochs=0, eval_use_exact_rotated_iou=True, eval_use_exact_rotated_iou_for_final_map=None)

Complete training loop.

Parameters:

Name	Type	Description	Default
model	Module	Model to train	required
train_loader	DataLoader	DataLoader for training	required
optimizer	Optimizer	Optimizer	required
device	device	Device to train on	required
num_epochs	int	Number of epochs to train	10
val_loader	Optional[DataLoader]	Optional validation DataLoader	None
lr_scheduler	Optional[Any]	Optional learning rate scheduler	None
checkpoint_manager	Optional[CheckpointManager]	Optional checkpoint manager	None
print_freq	int	Frequency of printing metrics	10
gradient_accumulation_steps	int	Gradient accumulation steps	1
max_grad_norm	Optional[float]	Maximum gradient norm	None
use_amp	bool	Use automatic mixed precision	False
loss_weights	Optional[Union[Dict[str, float], Callable[[int], Optional[Dict[str, float]]]]]	Loss component weights	None
start_epoch	int	Starting epoch (for resuming)	0
writer	Optional[Any]	Optional TensorBoard SummaryWriter for logging	None
class_map	Optional[Dict[str, int]]	Optional mapping from class names to class IDs (for evaluation)	None
class_names	Optional[Sequence[str]]	Optional list of class names (for evaluation)	None
eval_score_threshold	float	Minimum confidence score for evaluation detections	0.05
eval_iou_threshold	float	IoU threshold for mAP calculation	0.5
eval_compute_map_final	bool	If True, after training load the best checkpoint and compute mAP once	True
eval_compute_map_every_n_epochs	int	If >0, compute mAP every N epochs during training (current model) (in addition to the final epoch). Use 0 to compute only at final epoch.	0
log_images	bool	Whether to log prediction images to TensorBoard	True
max_images_to_log	int	Maximum number of images to log per epoch	4
fixed_image_index	int	Index in validation order of the image to always log (0 = first image of first batch)	0
log_random_image	bool	If True, also log one randomly selected image from the pool	True
image_pool_max_size	int	Maximum size of the pool from which random images are drawn	128
log_debug_anchors_proposals	bool	If True, log anchors and RPN proposals to TensorBoard (random val image each epoch; for debugging).	False
profiler	Optional[Any]	Optional TrainingProfiler instance for performance profiling (should be used as context manager)	None
normalize_mean	Optional[list[float]]	Normalization mean values [R, G, B] for denormalization. Defaults to ImageNet mean.	None
normalize_std	Optional[list[float]]	Normalization std values [R, G, B] for denormalization. Defaults to ImageNet std.	None
vis_image_size	Optional[Tuple[int, int]]	Optional (height, width) for TensorBoard prediction images; when set, image is cropped to this size so boxes align (default (1024, 1024) when passed from train script).	None
train_sampler	Optional[Any]	Optional DistributedSampler for training (call set_epoch each epoch)	None
rank	Optional[int]	Optional rank for DDP (only rank 0 saves checkpoints and runs validation)	None
progress_stream	Optional[Any]	If set (e.g. original stderr when stdout is teed to a log file), tqdm writes here so the progress bar is not duplicated in the log file.	None
early_stop_patience	Optional[int]	If set, stop after this many epochs without improvement on early_stop_metric (skipped epochs where mAP was not computed, e.g. mAP=-1).	None
early_stop_metric	str	Validation metric key to monitor (default mAP).	'mAP'
early_stop_min_delta	float	Minimum relative change to count as improvement.	0.0
early_stop_higher_is_better	Optional[bool]	If None, inferred from metric name (mAP / AP_* -> True).	None
freeze_backbone_epochs	int	If >0, freeze backbone.* while epoch < this value (0-based).	0
freeze_rpn_epochs	int	If >0, freeze rpn_head.* while epoch < this value; no-op without RPN.	0

Returns:

Type	Description
Dict[str, Any]	Dictionary with training history

Source code in oriented_det/train/engine.py

def train(
    model: nn.Module,
    train_loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    *,
    num_epochs: int = 10,
    val_loader: Optional[DataLoader] = None,
    lr_scheduler: Optional[Any] = None,
    checkpoint_manager: Optional[CheckpointManager] = None,
    print_freq: int = 10,
    gradient_accumulation_steps: int = 1,
    max_grad_norm: Optional[float] = None,
    use_amp: bool = False,
    loss_weights: Optional[Union[Dict[str, float], Callable[[int], Optional[Dict[str, float]]]]] = None,
    roi_class_weights: Optional[Callable[[int], Optional[torch.Tensor]]] = None,
    start_epoch: int = 0,
    writer: Optional[Any] = None,
    class_map: Optional[Dict[str, int]] = None,
    class_names: Optional[Sequence[str]] = None,
    eval_score_threshold: float = 0.05,
    eval_per_class_score_threshold: Optional[Dict[str, float]] = None,
    eval_vis_score_threshold: Optional[float] = None,
    eval_iou_threshold: float = 0.5,
    eval_extended_gt_metrics: bool = False,
    eval_compute_map_final: bool = True,
    eval_compute_map_every_n_epochs: int = 0,
    log_images: bool = True,
    max_images_to_log: int = 4,
    fixed_image_index: int = 0,
    log_random_image: bool = True,
    image_pool_max_size: int = 128,
    log_debug_anchors_proposals: bool = False,
    profiler: Optional[Any] = None,
    normalize_mean: Optional[list[float]] = None,
    normalize_std: Optional[list[float]] = None,
    vis_image_size: Optional[Tuple[int, int]] = None,
    train_sampler: Optional[Any] = None,
    val_sampler: Optional[Any] = None,
    rank: Optional[int] = None,
    progress_stream: Optional[Any] = None,
    debug: bool = False,
    lr_scheduler_plateau_metric: str = "total_loss",
    early_stop_patience: Optional[int] = None,
    early_stop_metric: str = "mAP",
    early_stop_min_delta: float = 0.0,
    early_stop_higher_is_better: Optional[bool] = None,
    freeze_backbone_epochs: int = 0,
    freeze_rpn_epochs: int = 0,
    eval_use_exact_rotated_iou: bool = True,
    eval_use_exact_rotated_iou_for_final_map: Optional[bool] = None,
) -> Dict[str, Any]:
    """Complete training loop.

    Args:
        model: Model to train
        train_loader: DataLoader for training
        optimizer: Optimizer
        device: Device to train on
        num_epochs: Number of epochs to train
        val_loader: Optional validation DataLoader
        lr_scheduler: Optional learning rate scheduler
        checkpoint_manager: Optional checkpoint manager
        print_freq: Frequency of printing metrics
        gradient_accumulation_steps: Gradient accumulation steps
        max_grad_norm: Maximum gradient norm
        use_amp: Use automatic mixed precision
        loss_weights: Loss component weights
        start_epoch: Starting epoch (for resuming)
        writer: Optional TensorBoard SummaryWriter for logging
        class_map: Optional mapping from class names to class IDs (for evaluation)
        class_names: Optional list of class names (for evaluation)
        eval_score_threshold: Minimum confidence score for evaluation detections
        eval_iou_threshold: IoU threshold for mAP calculation
        eval_compute_map_final: If True, after training load the best checkpoint and compute mAP once
        eval_compute_map_every_n_epochs: If >0, compute mAP every N epochs during training (current model)
            (in addition to the final epoch). Use 0 to compute only at final epoch.
        log_images: Whether to log prediction images to TensorBoard
        max_images_to_log: Maximum number of images to log per epoch
        fixed_image_index: Index in validation order of the image to always log (0 = first image of first batch)
        log_random_image: If True, also log one randomly selected image from the pool
        image_pool_max_size: Maximum size of the pool from which random images are drawn
        log_debug_anchors_proposals: If True, log anchors and RPN proposals to TensorBoard (random val image each epoch; for debugging).
        profiler: Optional TrainingProfiler instance for performance profiling (should be used as context manager)
        normalize_mean: Normalization mean values [R, G, B] for denormalization. Defaults to ImageNet mean.
        normalize_std: Normalization std values [R, G, B] for denormalization. Defaults to ImageNet std.
        vis_image_size: Optional (height, width) for TensorBoard prediction images; when set, image is cropped to this size so boxes align (default (1024, 1024) when passed from train script).
        train_sampler: Optional DistributedSampler for training (call set_epoch each epoch)
        rank: Optional rank for DDP (only rank 0 saves checkpoints and runs validation)
        progress_stream: If set (e.g. original stderr when stdout is teed to a log file),
            tqdm writes here so the progress bar is not duplicated in the log file.
        early_stop_patience: If set, stop after this many epochs without improvement on early_stop_metric
            (skipped epochs where mAP was not computed, e.g. mAP=-1).
        early_stop_metric: Validation metric key to monitor (default mAP).
        early_stop_min_delta: Minimum relative change to count as improvement.
        early_stop_higher_is_better: If None, inferred from metric name (mAP / AP_* -> True).
        freeze_backbone_epochs: If >0, freeze ``backbone.*`` while ``epoch <`` this value (0-based).
        freeze_rpn_epochs: If >0, freeze ``rpn_head.*`` while ``epoch <`` this value; no-op without RPN.

    Returns:
        Dictionary with training history
    """
    _require_torch()

    train_tracker = MetricTracker()
    val_tracker = MetricTracker() if val_loader is not None else None

    history = {
        "train": [],
        "val": [],
        "early_stopped": False,
    }

    # Calculate global step offset based on start_epoch
    global_step = start_epoch * len(train_loader) if start_epoch > 0 else 0

    # Track previous epoch's validation metrics for delta computation
    previous_val_metrics: Optional[Dict[str, Any]] = None

    _es_patience = int(early_stop_patience) if early_stop_patience is not None else 0
    _es_metric_key = (early_stop_metric or "mAP").strip()
    _es_best: Optional[float] = None
    _es_epochs_no_improve = 0

    def _infer_early_stop_higher_is_better() -> bool:
        if early_stop_higher_is_better is not None:
            return bool(early_stop_higher_is_better)
        lk = _es_metric_key.lower()
        if lk == "map" or lk.startswith("ap_") or lk.startswith("ap"):
            return True
        if "loss" in lk:
            return False
        return True

    _es_higher = _infer_early_stop_higher_is_better()

    epoch_wall_times: List[float] = []
    train_loop_t0 = time.perf_counter()
    train_started_at_iso = datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")
    if rank is None or rank == 0:
        print(f"Training started at: {train_started_at_iso}")

    # After partial freeze (freeze_backbone_epochs / freeze_rpn_epochs), all parameters train; rebuild DDP
    # with find_unused_parameters=False to avoid extra autograd traversals and PyTorch warnings.
    ddp_rebuilt_after_partial_freeze = False

    for epoch in range(start_epoch, start_epoch + num_epochs):
        epoch_t0 = time.perf_counter()
        val_metrics: Dict[str, Any] = {}
        if rank is None or rank == 0:
            epoch_idx = (epoch - start_epoch + 1)
            print(f"\nEpoch {epoch_idx}/{num_epochs}")
            print("-" * 50)

        # Set epoch for DistributedSampler (required for proper shuffling across epochs)
        if train_sampler is not None:
            train_sampler.set_epoch(epoch)
        # Also set epoch for validation sampler if using DistributedSampler
        if val_sampler is not None:
            val_sampler.set_epoch(epoch)

        train_model = model.module if hasattr(model, 'module') else model
        _fb = int(freeze_backbone_epochs or 0)
        _fr = int(freeze_rpn_epochs or 0)
        if _fb > 0 or _fr > 0:
            freeze_bb = epoch < _fb
            freeze_rpn = epoch < _fr
            set_backbone_requires_grad(train_model, freeze=freeze_bb)
            set_rpn_requires_grad(train_model, freeze=freeze_rpn)
            if rank is None or rank == 0:
                if (
                    epoch == start_epoch
                    or (_fb > 0 and epoch == _fb)
                    or (_fr > 0 and epoch == _fr)
                ):
                    parts = []
                    if _fb > 0:
                        parts.append(
                            f"backbone {'FROZEN' if freeze_bb else 'trainable'} (freeze_backbone_epochs={_fb})"
                        )
                    if _fr > 0 and model_has_rpn_head(train_model):
                        parts.append(
                            f"RPN {'FROZEN' if freeze_rpn else 'trainable'} (freeze_rpn_epochs={_fr})"
                        )
                    if parts:
                        print(f"  {' | '.join(parts)} — loop epoch {epoch}")

        _freeze_span = max(_fb, _fr)
        if (
            not ddp_rebuilt_after_partial_freeze
            and _freeze_span > 0
            and epoch >= _freeze_span
            and dist is not None
            and dist.is_initialized()
            and isinstance(model, nn.parallel.DistributedDataParallel)
        ):
            inner = model.module
            model = nn.parallel.DistributedDataParallel(
                inner,
                device_ids=None,
                output_device=None,
                find_unused_parameters=False,
                broadcast_buffers=False,
            )
            ddp_rebuilt_after_partial_freeze = True
            if rank is None or rank == 0:
                print(
                    "  DDP: rebuilt with find_unused_parameters=False (partial freeze ended; faster reductions)."
                )

        # Update NMS IoU threshold from schedule (e.g. Rotated RetinaNet: lower = more suppression)
        if hasattr(train_model, "set_final_nms_iou_for_epoch"):
            train_model.set_final_nms_iou_for_epoch(epoch)
        if hasattr(train_model, "set_roi_box_reg_iou_weight_for_epoch"):
            train_model.set_roi_box_reg_iou_weight_for_epoch(epoch)
        elif hasattr(train_model, "set_box_reg_iou_weight_for_epoch"):
            train_model.set_box_reg_iou_weight_for_epoch(epoch)
        if hasattr(train_model, "set_roi_box_reg_angle_weight_for_epoch"):
            train_model.set_roi_box_reg_angle_weight_for_epoch(epoch)
        # Update ROI class weights from schedule (optional; supports ramp-up to avoid early classifier collapse).
        if roi_class_weights is not None:
            try:
                w = roi_class_weights(epoch)
            except TypeError:
                # Backward compat: tolerate a constant tensor passed accidentally.
                w = roi_class_weights  # type: ignore[assignment]
            if w is not None:
                w = w.to(device=device, dtype=torch.float32)
                if hasattr(train_model, "set_class_weights_tensor"):
                    train_model.set_class_weights_tensor(w)
                elif hasattr(train_model, "roi_class_weights_tensor"):
                    train_model.roi_class_weights_tensor = w  # type: ignore[attr-defined]

        if hasattr(train_model, "set_grouped_ce_alpha_for_epoch"):
            train_model.set_grouped_ce_alpha_for_epoch(epoch)

        # Get loss weights for this epoch (support callable for dynamic weights)
        epoch_loss_weights = None
        if loss_weights is not None:
            if callable(loss_weights):
                epoch_loss_weights = loss_weights(epoch)
            else:
                epoch_loss_weights = loss_weights

        # Train
        train_metrics = train_one_epoch(
            model=model,
            data_loader=train_loader,
            optimizer=optimizer,
            device=device,
            epoch=epoch,
            epoch_idx=(epoch - start_epoch + 1),
            epoch_count=num_epochs,
            print_freq=print_freq,
            gradient_accumulation_steps=gradient_accumulation_steps,
            max_grad_norm=max_grad_norm,
            use_amp=use_amp,
            metric_tracker=train_tracker,
            loss_weights=epoch_loss_weights,
            writer=writer,
            global_step=global_step,
            profiler=profiler,
            lr_scheduler=lr_scheduler,
            rank=rank,
            progress_stream=progress_stream,
            debug=debug,
        )
        # Update global step from metrics if available
        if writer is not None and "_global_step" in train_metrics:
            global_step = train_metrics.pop("_global_step")
        history["train"].append(train_metrics)

        # Synchronize all ranks before validation (ensure training is complete on all ranks)
        # CRITICAL: All ranks must reach this barrier before validation starts
        # This prevents DDP operations in the next training epoch from interfering with validation
        if rank is not None and dist is not None and dist.is_initialized():
            dist.barrier()

        # Validate: All ranks participate in forward pass (prevents DDP deadlock)
        # Only rank 0 collects/logs metrics to avoid duplication
        if val_loader is not None and val_tracker is not None:
            # Compute mAP every N epochs during training (optional). Final mAP uses best model after loop.
            epoch_idx = (epoch - start_epoch + 1)
            map_every_n = max(0, int(eval_compute_map_every_n_epochs or 0))
            is_periodic_map_epoch = map_every_n > 0 and (epoch_idx % map_every_n == 0)
            compute_map_this_epoch = is_periodic_map_epoch

            if compute_map_this_epoch and (rank is None or rank == 0):
                print(f"\nComputing mAP (periodic evaluation: every {map_every_n} epoch(s))...")
            # Keep expensive GT–IoU diagnostics on the same schedule as mAP by default.
            extended_gt_this_epoch = bool(eval_extended_gt_metrics) and bool(compute_map_this_epoch)
            matching_metrics_this_epoch = compute_map_this_epoch or extended_gt_this_epoch

            # All ranks run validation forward pass (prevents DDP deadlock)
            # Only rank 0 collects metrics and logs
            val_metrics = evaluate(
                model=model,
                data_loader=val_loader,
                device=device,
                metric_tracker=val_tracker if (rank is None or rank == 0) else None,
                writer=writer if (rank is None or rank == 0) else None,
                epoch=epoch,
                class_map=class_map,
                class_names=class_names,
                score_threshold=eval_score_threshold,
                per_class_score_threshold=eval_per_class_score_threshold,
                vis_score_threshold=eval_vis_score_threshold,
                iou_threshold=eval_iou_threshold,
                extended_gt_metrics=extended_gt_this_epoch,
                compute_map=compute_map_this_epoch and (rank is None or rank == 0),  # Only rank 0 computes mAP
                log_images=log_images and (rank is None or rank == 0),  # Only rank 0 logs images
                max_images_to_log=max_images_to_log,
                fixed_image_index=fixed_image_index,
                log_random_image=log_random_image,
                image_pool_max_size=image_pool_max_size,
                log_debug_anchors_proposals=log_debug_anchors_proposals,
                normalize_mean=normalize_mean,
                normalize_std=normalize_std,
                vis_image_size=vis_image_size,
                rank=rank,
                progress_stream=progress_stream,
                debug=debug,
                eval_use_exact_rotated_iou=eval_use_exact_rotated_iou,
                compute_matching_metrics=matching_metrics_this_epoch,
            )
            # Only rank 0 has meaningful metrics; other ranks return empty dict
            if rank is None or rank == 0:
                history["val"].append(val_metrics)
                print(_format_validation_metrics(val_metrics, previous_val_metrics))
                previous_val_metrics = _snapshot_val_metrics_for_comparison(
                    val_metrics, previous_val_metrics
                )
            else:
                # Non-zero ranks return empty metrics (they still ran forward pass for DDP sync)
                history["val"].append({})

        # Synchronize all DDP ranks after validation.
        #
        # IMPORTANT: rank 0 may spend a long time computing mAP while other ranks wait.
        # Using an NCCL barrier can trip the NCCL watchdog timeout and abort the job.
        # Prefer a Gloo barrier when available (we already create a Gloo group for gather_object).
        if rank is not None and dist is not None and dist.is_initialized():
            try:
                gather_group = _get_gloo_gather_object_group()
                if gather_group is not None:
                    dist.barrier(group=gather_group)
                else:
                    dist.barrier()
            except TypeError:
                # Older PyTorch: barrier may not accept group=
                dist.barrier()

        # Learning rate scheduling (per epoch for standard schedulers)
        # Note: Warmup schedulers are stepped per optimizer step in train_one_epoch
        if lr_scheduler is not None:
            # Handle warmup scheduler's base scheduler stepping per epoch
            if hasattr(lr_scheduler, 'step_epoch'):
                # Custom warmup scheduler - step base scheduler per epoch
                lr_scheduler.step_epoch()
            elif isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                if val_loader is not None and val_metrics:
                    metric_val = val_metrics.get(
                        lr_scheduler_plateau_metric,
                        train_metrics.get(lr_scheduler_plateau_metric, 0.0),
                    )
                    # Do not step on sentinel mAP=-1 when periodic mAP was skipped this epoch
                    if (
                        lr_scheduler_plateau_metric.strip().lower() == "map"
                        and isinstance(metric_val, (int, float))
                        and float(metric_val) == -1.0
                    ):
                        pass
                    else:
                        lr_scheduler.step(metric_val)
                else:
                    metric_val = train_metrics.get(lr_scheduler_plateau_metric, 0.0)
                    lr_scheduler.step(metric_val)
            else:
                # Standard PyTorch scheduler (StepLR, etc.) — step once per epoch.
                # PyTorch 2.x uses LRScheduler; older versions use _LRScheduler.
                _base = getattr(torch.optim.lr_scheduler, 'LRScheduler', None) or getattr(torch.optim.lr_scheduler, '_LRScheduler', None)
                if _base is not None and isinstance(lr_scheduler, _base):
                    lr_scheduler.step()
            # Custom warmup schedulers are already stepped per optimizer step in train_one_epoch

        # Checkpointing (only on rank 0 for DDP, or all ranks if rank is None)
        if checkpoint_manager is not None and (rank is None or rank == 0):
            checkpoint_manager.save(
                model=model,
                optimizer=optimizer,
                scheduler=lr_scheduler,
                epoch=epoch,
                metrics=train_metrics,
            )
            if val_loader is not None and (rank is None or rank == 0):
                # Use val_metrics if available, otherwise train_metrics
                metrics_for_best = (
                    val_metrics if (len(history["val"]) > 0 and history["val"][-1]) else train_metrics
                )
            else:
                # No validation pass: best checkpoint follows training metrics
                metrics_for_best = train_metrics
            checkpoint_manager.save_best(
                model=model,
                optimizer=optimizer,
                scheduler=lr_scheduler,
                epoch=epoch,
                metrics=metrics_for_best,
            )

        if rank is None or rank == 0:
            epoch_wall_s = time.perf_counter() - epoch_t0
            epoch_wall_times.append(epoch_wall_s)
            avg_s = sum(epoch_wall_times) / len(epoch_wall_times)
            remaining_epochs = start_epoch + num_epochs - epoch - 1
            map_n = max(0, int(eval_compute_map_every_n_epochs or 0))
            want_final_map = bool(eval_compute_map_final and val_loader is not None)
            parts = [
                f"  Timing: {_format_duration_hms(epoch_wall_s)} this epoch (avg {_format_duration_hms(avg_s)})",
            ]
            if remaining_epochs > 0:
                eta_total = avg_s * remaining_epochs + (avg_s if want_final_map else 0.0)
                second = (
                    f"ETA ~{_format_duration_hms(eta_total)} for {remaining_epochs} epoch(s) left"
                )
                if map_n > 0:
                    second += f", mAP every {map_n} epoch(s)"
                if want_final_map:
                    second += " and final mAP."
                else:
                    second += "."
                parts.append(second)
            elif want_final_map:
                parts.append(f"ETA ~{_format_duration_hms(avg_s)} for final mAP on best checkpoint.")
            else:
                parts.append("Training loop finished (no final mAP).")
            print(" | ".join(parts))

        # Early stopping on validation metric (rank 0 decides; broadcast to all DDP ranks)
        stop_flag = torch.zeros(1, dtype=torch.int32, device=device)
        if _es_patience > 0 and val_loader is not None and (rank is None or rank == 0):
            raw = val_metrics.get(_es_metric_key) if val_metrics else None
            valid = False
            v = 0.0
            if raw is not None:
                try:
                    v = float(raw)
                except (TypeError, ValueError):
                    valid = False
                else:
                    valid = not math.isnan(v) and not (_es_metric_key.lower() == "map" and v == -1.0)
            if valid:
                if _es_best is None:
                    _es_best = v
                    _es_epochs_no_improve = 0
                else:
                    improved = (
                        (v > _es_best + early_stop_min_delta)
                        if _es_higher
                        else (v < _es_best - early_stop_min_delta)
                    )
                    if improved:
                        _es_best = v
                        _es_epochs_no_improve = 0
                    else:
                        _es_epochs_no_improve += 1
                if _es_epochs_no_improve >= _es_patience:
                    stop_flag[0] = 1
                    history["early_stopped"] = True
                    history["early_stop_epoch"] = int(epoch)
                    history["early_stop_best_metric"] = float(_es_best) if _es_best is not None else None
                    print(
                        f"\nEarly stopping: {_es_metric_key} did not improve for {_es_patience} "
                        f"epoch(s) (best={_es_best:.6f})."
                    )
        if rank is not None and dist is not None and dist.is_initialized():
            dist.broadcast(stop_flag, src=0)
        if stop_flag.item() != 0:
            break

        # Free unused GPU memory between epochs to reduce OOM risk from fragmentation
        if device.type == "cuda":
            torch.cuda.empty_cache()

    # Final mAP with best checkpoint (after training loop)
    # DDP: all ranks must call evaluate() because evaluate() uses dist.gather_object on the val set.
    # If only rank 0 runs evaluate(), gather_object waits forever for other ranks → hung job, no mAP in log.
    _ran_final_map_evaluate = False
    if eval_compute_map_final and val_loader is not None:
        best_path = checkpoint_manager.best_checkpoint_path if checkpoint_manager is not None else None
        # In DDP, only rank 0 updates checkpoint_manager.best_checkpoint_path during save_best().
        # Broadcast the resolved best checkpoint path so every rank enters final evaluate()
        # and participates in the validation collectives.
        if rank is not None and dist is not None and dist.is_initialized():
            best_path_str = str(best_path) if best_path is not None else None
            best_path_list = [best_path_str] if rank == 0 else [None]
            dist.broadcast_object_list(best_path_list, src=0)
            best_path = Path(best_path_list[0]) if best_path_list[0] else None
        if best_path is not None and best_path.exists():
            _ran_final_map_evaluate = True
            _final_map_exact_iou = (
                eval_use_exact_rotated_iou
                if eval_use_exact_rotated_iou_for_final_map is None
                else bool(eval_use_exact_rotated_iou_for_final_map)
            )
            if rank is None or rank == 0:
                print("\nComputing final mAP using best checkpoint...")
                print(
                    f"  Final mAP IoU backend: "
                    f"{'exact CPU polygon' if _final_map_exact_iou else 'GPU sampling (approx)'}"
                )
            ckpt = torch.load(best_path, map_location=device)
            model_to_load = model.module if hasattr(model, "module") else model
            state_dict = ckpt["model_state_dict"]
            # Strip DDP "module." prefix so keys match the unwrapped model (same machine, saved under DDP)
            ckpt_has_module = any(k.startswith("module.") for k in state_dict.keys())
            model_has_module = any(k.startswith("module.") for k in model_to_load.state_dict().keys())
            if ckpt_has_module and not model_has_module:
                state_dict = {(k[7:] if k.startswith("module.") else k): v for k, v in state_dict.items()}
            elif model_has_module and not ckpt_has_module:
                state_dict = {f"module.{k}": v for k, v in state_dict.items()}
            model_to_load.load_state_dict(state_dict, strict=True)
            final_metrics = evaluate(
                model=model,
                data_loader=val_loader,
                device=device,
                metric_tracker=val_tracker if (rank is None or rank == 0) else None,
                writer=writer if (rank is None or rank == 0) else None,
                epoch=start_epoch + num_epochs - 1,
                class_map=class_map,
                class_names=class_names,
                score_threshold=eval_score_threshold,
                per_class_score_threshold=eval_per_class_score_threshold,
                vis_score_threshold=eval_vis_score_threshold,
                iou_threshold=eval_iou_threshold,
                extended_gt_metrics=eval_extended_gt_metrics,
                compute_map=(rank is None or rank == 0),
                log_images=False,
                max_images_to_log=0,
                fixed_image_index=fixed_image_index,
                log_random_image=False,
                image_pool_max_size=image_pool_max_size,
                log_debug_anchors_proposals=False,
                normalize_mean=normalize_mean,
                normalize_std=normalize_std,
                vis_image_size=vis_image_size,
                rank=rank,
                progress_stream=progress_stream,
                debug=debug,
                eval_use_exact_rotated_iou=_final_map_exact_iou,
                compute_matching_metrics=True,
            )
            if rank is None or rank == 0:
                print(_format_validation_metrics(final_metrics))
        elif rank is None or rank == 0:
            print("\nSkipping final mAP: no best checkpoint saved.")

    # DDP: rank 0 runs compute_oriented_map inside evaluate() (minutes); other ranks skip mAP and return
    # from evaluate() first. Without this barrier they exit train() and destroy_process_group() while
    # rank 0 is still in mAP → NCCL ALLREDUCE watchdog timeout (default ~30 min) on shutdown.
    if (
        _ran_final_map_evaluate
        and rank is not None
        and dist is not None
        and dist.is_initialized()
    ):
        dist.barrier()

    # Close TensorBoard writer if provided
    if writer is not None:
        writer.close()

    finished_at_iso = datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")
    total_wall_s = time.perf_counter() - train_loop_t0
    timing: Dict[str, Any] = {
        "started_at": train_started_at_iso,
        "finished_at": finished_at_iso,
        "total_wall_seconds": float(total_wall_s),
    }
    if epoch_wall_times:
        timing["epoch_wall_times_seconds"] = list(epoch_wall_times)
        timing["epochs_completed"] = len(epoch_wall_times)
        timing["mean_epoch_seconds"] = float(sum(epoch_wall_times) / len(epoch_wall_times))
        timing["min_epoch_seconds"] = float(min(epoch_wall_times))
        timing["max_epoch_seconds"] = float(max(epoch_wall_times))
    history["timing"] = timing
    if rank is None or rank == 0:
        print("\n" + "=" * 80)
        print("Training timing summary")
        print("=" * 80)
        print(f"  Started:   {train_started_at_iso}")
        print(f"  Finished:  {finished_at_iso}")
        print(f"  Total wall: {_format_duration_hms(total_wall_s)} ({total_wall_s:.1f} s)")
        if epoch_wall_times:
            print(
                f"  Epochs:    {len(epoch_wall_times)} completed "
                f"(mean {_format_duration_hms(timing['mean_epoch_seconds'])}, "
                f"min {_format_duration_hms(timing['min_epoch_seconds'])}, "
                f"max {_format_duration_hms(timing['max_epoch_seconds'])})"
            )
        else:
            print("  Epochs:    no epoch timing recorded (0 loop iterations).")
        if history.get("early_stopped"):
            print("  Note:      training stopped early (early stopping).")
        print("=" * 80)

    return history

train_one_epoch(model, data_loader, optimizer, device, *, epoch=0, epoch_idx=None, epoch_count=None, print_freq=10, gradient_accumulation_steps=1, max_grad_norm=None, use_amp=False, metric_tracker=None, loss_weights=None, writer=None, global_step=None, profiler=None, lr_scheduler=None, rank=None, progress_stream=None, debug=False)

Train model for one epoch.

Parameters:

Name	Type	Description	Default
model	Module	Model to train	required
data_loader	DataLoader	DataLoader for training data	required
optimizer	Optimizer	Optimizer	required
device	device	Device to train on	required
epoch	int	Current epoch number	0
print_freq	int	Frequency of printing metrics	10
gradient_accumulation_steps	int	Number of steps to accumulate gradients	1
max_grad_norm	Optional[float]	Maximum gradient norm for clipping	None
use_amp	bool	Use automatic mixed precision	False
metric_tracker	Optional[MetricTracker]	Optional metric tracker	None
loss_weights	Optional[Dict[str, float]]	Optional weights for different loss components	None
writer	Optional[Any]	Optional TensorBoard SummaryWriter for logging	None
global_step	Optional[int]	Optional global step counter (will be incremented internally)	None
profiler	Optional[Any]	Optional TrainingProfiler instance for performance profiling	None
rank	Optional[int]	Optional rank for DDP (only rank 0 logs/prints, None for single-GPU)	None
progress_stream	Optional[Any]	If set (e.g. original stderr when stdout is teed to a log file), tqdm writes here so the progress bar is not duplicated in the log file.	None

Returns:

Type	Description
Dict[str, float]	Dictionary of average metrics for the epoch

Source code in oriented_det/train/engine.py

def train_one_epoch(
    model: nn.Module,
    data_loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    *,
    epoch: int = 0,
    epoch_idx: Optional[int] = None,
    epoch_count: Optional[int] = None,
    print_freq: int = 10,
    gradient_accumulation_steps: int = 1,
    max_grad_norm: Optional[float] = None,
    use_amp: bool = False,
    metric_tracker: Optional[MetricTracker] = None,
    loss_weights: Optional[Dict[str, float]] = None,
    writer: Optional[Any] = None,
    global_step: Optional[int] = None,
    profiler: Optional[Any] = None,
    lr_scheduler: Optional[Any] = None,
    rank: Optional[int] = None,
    progress_stream: Optional[Any] = None,
    debug: bool = False,
) -> Dict[str, float]:
    """Train model for one epoch.

    Args:
        model: Model to train
        data_loader: DataLoader for training data
        optimizer: Optimizer
        device: Device to train on
        epoch: Current epoch number
        print_freq: Frequency of printing metrics
        gradient_accumulation_steps: Number of steps to accumulate gradients
        max_grad_norm: Maximum gradient norm for clipping
        use_amp: Use automatic mixed precision
        metric_tracker: Optional metric tracker
        loss_weights: Optional weights for different loss components
        writer: Optional TensorBoard SummaryWriter for logging
        global_step: Optional global step counter (will be incremented internally)
        profiler: Optional TrainingProfiler instance for performance profiling
        rank: Optional rank for DDP (only rank 0 logs/prints, None for single-GPU)
        progress_stream: If set (e.g. original stderr when stdout is teed to a log file),
            tqdm writes here so the progress bar is not duplicated in the log file.

    Returns:
        Dictionary of average metrics for the epoch
    """
    _require_torch()

    model.train()
    metric_tracker = metric_tracker or MetricTracker()
    # Get device type from device parameter (e.g., 'cuda' or 'cpu')
    device_type = device.type if hasattr(device, 'type') else str(device).split(':')[0]
    # Initialize GradScaler - check which API is available
    if use_amp and GradScaler is not None:
        # Check if we're using the new torch.amp API (PyTorch 2.3+) or old torch.cuda.amp API
        # In PyTorch 2.2.2, torch.amp.GradScaler doesn't exist, so we use torch.cuda.amp.GradScaler
        # which doesn't take a device parameter (CUDA is implicit)
        try:
            # Try to import the new API to see if it exists
            from torch.amp import GradScaler as AmpGradScaler
            # New API (PyTorch 2.3+): GradScaler(device) - device is first positional arg
            scaler = AmpGradScaler(device_type)
        except ImportError:
            # Old API (torch.cuda.amp.GradScaler, PyTorch 2.2.2): no device parameter, CUDA is implicit
            # Only use old API if CUDA is available, otherwise skip scaler
            if device_type == 'cuda':
                scaler = GradScaler()
            else:
                scaler = None
    else:
        scaler = None

    optimizer.zero_grad()

    num_batches = len(data_loader)
    total_loss = 0.0
    current_step = global_step if global_step is not None else 0
    skip_optimizer_step = False  # Set when we skip backward due to nan/inf loss

    # Create progress bar (only on rank 0 or when rank is None).
    # When progress_stream is set (e.g. original stderr), tqdm writes there so the log file doesn't get every update.
    import sys
    pbar = None
    if tqdm is not None and (rank is None or rank == 0):
        tqdm_file = progress_stream if progress_stream is not None else sys.stderr
        if epoch_idx is not None and epoch_count is not None:
            epoch_desc = f"Epoch {int(epoch_idx)}/{int(epoch_count)}"
        else:
            epoch_desc = f"Epoch {epoch + 1}"
        pbar = tqdm(
            total=num_batches,
            desc=epoch_desc,
            unit="batch",
            ncols=120,
            leave=True,
            file=tqdm_file,
            disable=False,
        )
        if progress_stream is None:
            sys.stdout.flush()


    for batch_idx, batch in enumerate(data_loader):
        step_start = time.time()

        try:
            # Unpack batch (format depends on collate function)
            if isinstance(batch, (list, tuple)) and len(batch) == 2:
                images, targets = batch
            else:
                images, targets = batch.get("images", []), batch.get("targets", [])

            # Move to device
            if isinstance(images, list):
                images = [img.to(device) if torch.is_tensor(img) else img for img in images]
            elif torch.is_tensor(images):
                images = images.to(device)

            # Move targets to device (tensors in target dicts)
            if isinstance(targets, list):
                for target in targets:
                    if isinstance(target, dict):
                        for key, value in target.items():
                            if torch.is_tensor(value):
                                target[key] = value.to(device)

            # Debug: first-batch GT stats (per-image and per-class counts) to spot data/annotation issues.
            if debug and batch_idx == 0 and (rank is None or rank == 0) and isinstance(targets, list):
                from collections import Counter
                gt_per_image = []
                class_counts = Counter()
                for t in targets:
                    if isinstance(t, dict) and "labels" in t:
                        labels = t["labels"]
                        n = len(labels) if torch.is_tensor(labels) else len(labels)
                        gt_per_image.append(n)
                        if torch.is_tensor(labels):
                            for c in labels.cpu().tolist():
                                class_counts[c] += 1
                        else:
                            for c in labels:
                                class_counts[c] += 1
                if gt_per_image:
                    print(
                        f"  [debug] First batch GT: images={len(gt_per_image)}, "
                        f"GT/image min={min(gt_per_image)}, max={max(gt_per_image)}, mean={sum(gt_per_image)/len(gt_per_image):.1f}; "
                        f"per-class counts: {dict(class_counts.most_common(10))}",
                        flush=True,
                    )

            # Forward pass
            if use_amp and autocast is not None:
                try:
                    # PyTorch 2.x: device_type required - use device type from device parameter
                    with autocast(device_type=device_type):
                        loss_dict = model(images, targets)
                except TypeError:
                    # Older PyTorch (torch.cuda.amp.autocast): no device_type parameter, CUDA is implicit
                    # Old API only accepts keyword arguments: enabled, dtype, cache_enabled
                    # Do not pass device_type as positional argument - it would be interpreted as enabled='cuda' (wrong!)
                    with autocast():
                        loss_dict = model(images, targets)
            else:
                loss_dict = model(images, targets)

            # Extract and weight losses
            if not isinstance(loss_dict, dict):
                raise ValueError("Model forward must return a dict of losses.")

            total_loss_value = 0.0
            for key, value in loss_dict.items():
                if not torch.is_tensor(value):
                    continue
                weight = 1.0 if loss_weights is None else loss_weights.get(key, 1.0)
                total_loss_value += weight * value

            # Normalize by accumulation steps
            total_loss_value = total_loss_value / gradient_accumulation_steps

            # Robustness: skip backward and optimizer step when loss is nan/inf to avoid corrupting weights
            loss_is_invalid = not (torch.is_tensor(total_loss_value) and torch.isfinite(total_loss_value).all())
            if loss_is_invalid:
                skip_optimizer_step = True
                if rank is None or rank == 0:
                    import warnings
                    warnings.warn(
                        f"Train batch {batch_idx}: loss is nan or inf (total_loss_value={total_loss_value}), "
                        "skipping backward and optimizer step for this accumulation window."
                    )
                # Use zero for logging so epoch average is not skewed
                total_loss_value_safe = total_loss_value
                if torch.is_tensor(total_loss_value):
                    total_loss_value_safe = torch.tensor(0.0, device=total_loss_value.device, dtype=total_loss_value.dtype)
            else:
                # Backward pass
                if use_amp and scaler is not None:
                    scaler.scale(total_loss_value).backward()
                else:
                    total_loss_value.backward()
                total_loss_value_safe = total_loss_value

            # Gradient accumulation
            if (batch_idx + 1) % gradient_accumulation_steps == 0:
                if skip_optimizer_step:
                    optimizer.zero_grad()
                    skip_optimizer_step = False
                    grad_norm = None
                else:
                    # Gradient clipping
                    grad_norm = None
                    if max_grad_norm is not None:
                        if use_amp and scaler is not None:
                            scaler.unscale_(optimizer)
                        # Compute gradient norm before clipping
                        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
                    else:
                        # Compute gradient norm even if not clipping
                        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), float('inf'))

                    # Optimizer step
                    if use_amp and scaler is not None:
                        scaler.step(optimizer)
                        scaler.update()
                    else:
                        optimizer.step()

                    optimizer.zero_grad()

                # Advance profiler step if provided
                if profiler is not None:
                    profiler.step()
                # Increment step counter and log to TensorBoard (per optimizer step)
                current_step += 1

                # Step learning rate scheduler only for WarmupScheduler (per optimizer step).
                # Do not step StepLR/etc. here — they are stepped once per epoch below.
                if lr_scheduler is not None and hasattr(lr_scheduler, 'step_epoch'):
                    lr_scheduler.step()

                if writer is not None:
                    # Update metrics for logging (use safe value when loss was invalid)
                    metrics = {k: float(v.item()) if torch.is_tensor(v) and torch.isfinite(v).all() else float(v) if not torch.is_tensor(v) else 0.0 for k, v in loss_dict.items()}
                    metrics["total_loss"] = float(total_loss_value_safe.item() * gradient_accumulation_steps)
                    for key, value in metrics.items():
                        # ROI diagnostics: log to TensorBoard only when debug is enabled (otherwise stdout/train.log only).
                        if key.startswith("roi_") and not debug:
                            continue
                        writer.add_scalar(f"train/{key}", value, current_step)
                    # Log learning rate (config-scale base, not backbone-only when param groups are used)
                    current_lr = _optimizer_reference_learning_rate(optimizer)
                    writer.add_scalar("train/learning_rate", current_lr, current_step)
                    # Log gradient norm
                    if grad_norm is not None:
                        writer.add_scalar("train/grad_norm", float(grad_norm), current_step)

                    # Log per-component gradient norms (TensorBoard)
                    if grad_norm is not None:
                        # Get the actual model (unwrap DDP if needed)
                        actual_model = model.module if hasattr(model, 'module') else model

                        # Compute gradient norms per component
                        component_grad_norms = {}
                        for name, param in actual_model.named_parameters():
                            if param.grad is not None:
                                param_grad_norm = param.grad.data.norm(2).item()
                                # Group by component (backbone, rpn, roi)
                                if 'backbone' in name:
                                    component = 'backbone'
                                elif 'rpn' in name.lower():
                                    component = 'rpn'
                                elif 'roi' in name.lower() or 'head' in name.lower():
                                    component = 'roi'
                                else:
                                    component = 'other'

                                if component not in component_grad_norms:
                                    component_grad_norms[component] = []
                                component_grad_norms[component].append(param_grad_norm)

                        # Log aggregated gradient norms per component
                        for component, norms in component_grad_norms.items():
                            if norms:
                                avg_norm = sum(norms) / len(norms)
                                max_norm = max(norms)
                                writer.add_scalar(f"train/grad_norm_{component}_avg", avg_norm, current_step)
                                writer.add_scalar(f"train/grad_norm_{component}_max", max_norm, current_step)

                        # Also log per-component total gradient norm (sum of squares)
                        for component, norms in component_grad_norms.items():
                            if norms:
                                total_norm = math.sqrt(sum(n * n for n in norms))
                                writer.add_scalar(f"train/grad_norm_{component}_total", total_norm, current_step)

            # Update metrics (for MetricTracker; use safe values when loss was invalid)
            metrics = {}
            for k, v in loss_dict.items():
                if torch.is_tensor(v):
                    metrics[k] = float(v.item()) if torch.isfinite(v).all() else 0.0
                else:
                    metrics[k] = float(v)
            metrics["total_loss"] = float(total_loss_value_safe.item() * gradient_accumulation_steps)
            metric_tracker.update(metrics)

            total_loss += metrics["total_loss"]

            # Update progress bar
            elapsed = time.time() - step_start
            metric_tracker.update_time(elapsed)

            # Update progress bar every batch (to avoid duplicate refreshes)
            if pbar is not None and (rank is None or rank == 0):
                # Compute summary for postfix (use rolling window, min of print_freq or current batch)
                window = min(print_freq, batch_idx + 1)
                summary = metric_tracker.get_summary(window=window)
                loss_val = summary.get('total_loss', 0.0)
                avg_time = summary.get('time_per_step', elapsed)

                # Calculate speed (batches per second) and ETA
                batches_per_sec = 1.0 / avg_time if avg_time > 0 else 0.0
                remaining_batches = num_batches - (batch_idx + 1)
                eta_seconds = remaining_batches * avg_time if avg_time > 0 else 0.0

                # Format ETA as HH:MM:SS or MM:SS
                if eta_seconds >= 3600:
                    eta_str = f"{int(eta_seconds // 3600):d}:{int((eta_seconds % 3600) // 60):02d}:{int(eta_seconds % 60):02d}"
                else:
                    eta_str = f"{int(eta_seconds // 60):d}:{int(eta_seconds % 60):02d}"

                # Update postfix and position together (single refresh)
                loss_str = f"Loss: {loss_val:.4f}"
                speed_str = f"{batches_per_sec:.2f}it/s"
                eta_str_tqdm = f"ETA: {eta_str}"
                pbar.set_postfix_str(f"{loss_str} | {speed_str} | {eta_str_tqdm}")
                pbar.update(1)  # This will refresh with the updated postfix

            # Print progress (every print_freq batches or at the end) - only for non-tqdm case
            if (batch_idx + 1) % print_freq == 0 or (batch_idx + 1) == num_batches:
                if rank is None or rank == 0:
                    if pbar is None:  # Only print if tqdm is not available
                        summary = metric_tracker.get_summary(window=print_freq)
                        loss_val = summary.get('total_loss', 0.0)
                        avg_time = summary.get('time_per_step', elapsed)

                        # Calculate speed (batches per second) and ETA
                        batches_per_sec = 1.0 / avg_time if avg_time > 0 else 0.0
                        remaining_batches = num_batches - (batch_idx + 1)
                        eta_seconds = remaining_batches * avg_time if avg_time > 0 else 0.0

                        # Format ETA as HH:MM:SS or MM:SS
                        if eta_seconds >= 3600:
                            eta_str = f"{int(eta_seconds // 3600):d}:{int((eta_seconds % 3600) // 60):02d}:{int(eta_seconds % 60):02d}"
                        else:
                            eta_str = f"{int(eta_seconds // 60):d}:{int(eta_seconds % 60):02d}"

                        disp_epoch = int(epoch_idx) if epoch_idx is not None else int(epoch + 1)
                        disp_total = int(epoch_count) if epoch_count is not None else None
                        if disp_total is not None:
                            prefix = f"Epoch {disp_epoch}/{disp_total}"
                        else:
                            prefix = f"Epoch {disp_epoch}"
                        print(
                            f"{prefix} [{batch_idx + 1}/{num_batches}] "
                            f"Loss: {loss_val:.4f} | "
                            f"Speed: {batches_per_sec:.2f}it/s | "
                            f"ETA: {eta_str}"
                        )
                    # Force flush to ensure output is visible
                    import sys
                    sys.stdout.flush()

        except Exception as e:
            if pbar is not None:
                pbar.close()
            if rank is None or rank == 0:
                print(f"Error in batch {batch_idx}: {e}")
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            raise

    # Flush any partial accumulation window at epoch end so tail microbatches are not dropped.
    remainder_batches = num_batches % gradient_accumulation_steps
    if num_batches > 0 and remainder_batches != 0:
        if skip_optimizer_step:
            optimizer.zero_grad()
            skip_optimizer_step = False
            grad_norm = None
        else:
            grad_norm = None
            if max_grad_norm is not None:
                if use_amp and scaler is not None:
                    scaler.unscale_(optimizer)
                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            else:
                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), float('inf'))

            if use_amp and scaler is not None:
                scaler.step(optimizer)
                scaler.update()
            else:
                optimizer.step()

            optimizer.zero_grad()

        if profiler is not None:
            profiler.step()
        current_step += 1

        if lr_scheduler is not None and hasattr(lr_scheduler, 'step_epoch'):
            lr_scheduler.step()

        if writer is not None:
            metrics = {
                k: float(v.item()) if torch.is_tensor(v) and torch.isfinite(v).all() else float(v) if not torch.is_tensor(v) else 0.0
                for k, v in loss_dict.items()
            }
            metrics["total_loss"] = float(total_loss_value_safe.item() * gradient_accumulation_steps)
            for key, value in metrics.items():
                if key.startswith("roi_") and not debug:
                    continue
                writer.add_scalar(f"train/{key}", value, current_step)
            current_lr = _optimizer_reference_learning_rate(optimizer)
            writer.add_scalar("train/learning_rate", current_lr, current_step)
            if grad_norm is not None:
                writer.add_scalar("train/grad_norm", float(grad_norm), current_step)

        if rank is None or rank == 0:
            print(
                f"  Flushed final accumulation window with {remainder_batches} batch(es) at epoch end.",
                flush=True,
            )

    # Close progress bar
    if pbar is not None:
        pbar.close()

    # Final metrics
    avg_metrics = metric_tracker.get_summary()
    avg_metrics["epoch"] = epoch

    # Single line to stdout (so log file gets one line per epoch when stdout is teed)
    if rank is None or rank == 0:
        loss_val = avg_metrics.get("total_loss", 0.0)
        disp_epoch = int(epoch_idx) if epoch_idx is not None else int(epoch + 1)
        disp_total = int(epoch_count) if epoch_count is not None else None
        if disp_total is not None:
            print(
                f"Epoch {disp_epoch}/{disp_total} complete: {num_batches}/{num_batches} batches, "
                f"loss: {loss_val:.4f}",
                flush=True,
            )
        else:
            print(
                f"Epoch {disp_epoch} complete: {num_batches}/{num_batches} batches, loss: {loss_val:.4f}",
                flush=True,
            )
        print(f"  Effective LR: {_format_effective_lrs(optimizer, lr_scheduler)}", flush=True)
        # ROI assignment diagnostics (Rotated Faster R-CNN: RPN proposals only, before add_gt_as_proposals).
        if "roi_num_pos" in avg_metrics:
            print(
                "  ROI hints: "
                f"pos={avg_metrics.get('roi_num_pos', 0.0):.1f}, "
                f"bg={avg_metrics.get('roi_num_bg', 0.0):.1f}, "
                f"ignore={avg_metrics.get('roi_num_ignore', 0.0):.1f}, "
                f"sampled_fg={avg_metrics.get('roi_sampled_fg', 0.0):.1f}, "
                f"sampled_bg={avg_metrics.get('roi_sampled_bg', 0.0):.1f}, "
                f"matched_gt={avg_metrics.get('roi_matched_gt', 0.0):.1f}, "
                f"match_rate={avg_metrics.get('roi_match_rate', 0.0):.2%}",
                flush=True,
            )
        # Debug: full loss breakdown to compare with MMRotate and spot unbalanced losses.
        if debug:
            loss_keys = [k for k in avg_metrics if k not in ("epoch", "_global_step") and (k == "total_loss" or k.startswith("loss_"))]
            if loss_keys:
                parts = [f"{k}={avg_metrics[k]:.4f}" for k in sorted(loss_keys)]
                print("  [debug] Loss breakdown: " + ", ".join(parts), flush=True)
            if writer is not None:
                current_lr = _optimizer_reference_learning_rate(optimizer)
                print(f"  [debug] Learning rate: {current_lr:.6e}", flush=True)

    # Log epoch-level metrics to TensorBoard
    if writer is not None:
        for key, value in avg_metrics.items():
            if key != "epoch":
                # ROI diagnostics: log to TensorBoard when debug is enabled.
                if key.startswith("roi_") and not debug:
                    continue
                writer.add_scalar(f"train_epoch/{key}", value, epoch)
        # Store updated step in metrics for retrieval by train() function
        avg_metrics["_global_step"] = current_step

    return avg_metrics

Important Notes

Training Features

The training engine provides efficient features:

• Mixed Precision Training: Automatic FP16 support via use_amp=True
• Gradient Accumulation: Train with large effective batch sizes
• Checkpointing: Automatic saving with best model tracking
• Metric Tracking: Built-in metric aggregation and reporting
• TensorBoard Logging: Optional TensorBoard integration for visualization
• Robust Error Handling: Graceful recovery from batch errors

Loss Components

For two-stage detectors like OrientedRCNN, there are four main loss components:

1. loss_objectness (RPN): Binary classification (object vs. background)

2. Expected: 0.5-0.7 early training → 0.1-0.3 converged

3. loss_rpn_box_reg (RPN): Box regression for refining anchors into proposals

4. Expected: 0.8-1.5 early training → 0.3-0.6 converged

5. loss_classifier (ROI): Multi-class classification for object classes

6. Expected: 0.5-1.0 early training → 0.2-0.4 converged

7. loss_box_reg (ROI): Box regression for refining proposals into final boxes

8. Expected: 0.8-1.5 early training → 0.2-0.5 converged

Memory Optimization

For memory-efficient training with OrientedRCNN:

• Use roi_chunk_size parameter to process ROIs in chunks (default: 32)
• Enable roi_use_checkpoint for gradient checkpointing (~2x less memory, ~30% slower)
• Recommended: roi_chunk_size=16 with checkpointing for 8-16GB GPUs

Examples

Basic Training

from oriented_det.train import train_one_epoch, train
from oriented_det.train import MetricTracker, CheckpointManager

# Single epoch
metrics = train_one_epoch(
    model=model,
    data_loader=train_loader,
    optimizer=optimizer,
    device=device,
    use_amp=True,  # Automatic mixed precision
    gradient_accumulation_steps=4,
)

# Full training loop
checkpoint_manager = CheckpointManager(
    "checkpoints/",
    best_metric="mAP",
    higher_is_better=True
)

history = train(
    model=model,
    train_loader=train_loader,
    optimizer=optimizer,
    device=device,
    num_epochs=12,
    val_loader=val_loader,
    checkpoint_manager=checkpoint_manager,
    use_amp=True,
    max_grad_norm=1.0,  # Gradient clipping
)

TensorBoard Logging

from torch.utils.tensorboard import SummaryWriter
from oriented_det.train import train

# Create TensorBoard writer
writer = SummaryWriter(log_dir="runs/experiment_1")

# Train with TensorBoard logging
history = train(
    model=model,
    train_loader=train_loader,
    optimizer=optimizer,
    device=device,
    num_epochs=12,
    val_loader=val_loader,
    writer=writer,  # Enable TensorBoard logging
)

# View logs with: tensorboard --logdir runs/experiment_1

Learning rate schedulers

The training loop accepts any PyTorch-style LR scheduler. When using ReduceLROnPlateau, pass lr_scheduler_plateau_metric to train() so the engine calls scheduler.step(metric_value) with the correct validation metric (e.g. "total_loss" or "mAP").

When using tools/train.py with a JSON config, set training.lr_scheduler_type to one of:

• multistep / step (default) — MultiStepLR or StepLR; optional warmup via lr_warmup_steps.
• reduce_on_plateau — ReduceLROnPlateau; configure lr_scheduler_plateau_metric, lr_scheduler_plateau_factor, lr_scheduler_plateau_patience.
• one_cycle — OneCycleLR (stepped every optimizer step); configure lr_scheduler_one_cycle_* options.
• cosine_annealing / cosine — PyTorch CosineAnnealingLR (lr_scheduler_cosine_t_max / lr_scheduler_cosine_epochs).
• cosine_annealing_with_tail / cosine_with_tail — Cosine then fixed lr_scheduler_cosine_tail_lr.

See the Training user guide — Learning rate scheduling for config examples and usage.

Performance Profiling

from oriented_det.train.profiler import TrainingProfiler
from torch.profiler import ProfilerActivity, schedule

# Create profiler
profiler = TrainingProfiler(
    log_dir="runs/profiling",
    activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
    schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
    record_shapes=True,
    profile_memory=True,
)

# Profile a few training steps
model.train()
with profiler:
    for i, (images, targets) in enumerate(train_loader):
        if i >= 5:  # Profile first 5 batches
            break
        # ... training code ...
        profiler.step()

# Print summary
profiler.print_summary(sort_by="cuda_time_total", row_limit=30)

# View detailed trace in TensorBoard
# tensorboard --logdir runs/profiling