Train a YOLOX Model

SUMMARY

Train YOLOX object detectors through the shared Iris dataset, Trainer, checkpoint, export, and inference APIs. YOLOXConfig controls the variant, COCO category mapping, augmentation, EMA, multiscale training, and optimization schedule.

Training Workflow

The example below reads category IDs from the COCO metadata, exposes model and Trainer settings as command-line arguments, creates training and validation datasets, and optionally resumes from a checkpoint.

Create train_yolox.py

"""Train a YOLOX detector on a COCO dataset."""

import argparse
from pathlib import Path

import torch
from loguru import logger

from telekinesis.iris.dataset import COCODataset, prepare_coco_datasets
from telekinesis.iris.logger import TensorBoardMetricLogger
from telekinesis.iris.models import YOLOX, YOLOXConfig, YOLOXVariant
from telekinesis.iris.optimizer import OptimizationConfig
from telekinesis.iris.trainer import Precision, Trainer


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--run-name", default="yolox")
    parser.add_argument(
        "--variant",
        choices=[variant.value for variant in YOLOXVariant],
        default=YOLOXVariant.SMALL.value,
    )
    parser.add_argument("--resolution", type=int)
    parser.add_argument("--pretrained", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--epochs", type=int, default=300)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--learning-rate", type=float, default=0.01)
    parser.add_argument("--weight-decay", type=float, default=5e-4)
    parser.add_argument("--nms-threshold", type=float, default=0.65)
    parser.add_argument("--max-detections", type=int, default=300)
    parser.add_argument("--no-aug-epochs", type=int, default=15)
    parser.add_argument("--warmup-epochs", type=float, default=5.0)
    parser.add_argument("--validation-split", type=float, default=0.3)
    parser.add_argument("--evaluation-interval", type=int, default=10)
    parser.add_argument("--checkpoint-interval", type=int, default=10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", help="For example: cuda, cuda:0, or cpu.")
    parser.add_argument(
        "--precision",
        choices=[precision.value for precision in Precision],
        default=Precision.FP32.value,
    )
    parser.add_argument("--resume-from", type=Path)
    return parser.parse_args()


def train_yolox(args: argparse.Namespace) -> None:
    """Configure and train YOLOX."""
    metadata_dir = args.dataset_dir / "train"
    if not metadata_dir.is_dir():
        metadata_dir = args.dataset_dir
    metadata = COCODataset(metadata_dir)
    torch.manual_seed(args.seed)

    config = YOLOXConfig(
        variant=YOLOXVariant(args.variant),
        num_classes=len(metadata.category_ids),
        category_ids=tuple(metadata.category_ids),
        pretrained=args.pretrained,
        resolution=args.resolution,
        nms_threshold=args.nms_threshold,
        max_detections=args.max_detections,
        no_aug_epochs=args.no_aug_epochs,
        warmup_epochs=args.warmup_epochs,
        optimization=OptimizationConfig(
            optimizer="sgd",
            learning_rate=args.learning_rate,
            weight_decay=args.weight_decay,
            optimizer_options={"momentum": 0.9, "nesterov": True},
            scheduler="cosine",
            scheduler_options={"min_factor": 0.05},
            scheduler_interval="step",
        ),
    )
    model = YOLOX(config=config)
    training_dataset, validation_dataset = prepare_coco_datasets(
        dataset_dir=args.dataset_dir,
        validation_split=args.validation_split,
        seed=args.seed,
        train_transforms=model.train_transforms,
        val_transforms=model.val_transforms,
        include_masks=model.requires_masks,
    )
    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Training YOLOX on {len(training_dataset)} images using {device}.")
    metric_logger = TensorBoardMetricLogger(args.output_dir / "tensorboard")
    model_trainer = Trainer(
        model=model,
        run_name=args.run_name,
        training_dataset=training_dataset,
        validation_dataset=validation_dataset,
        output_dir=args.output_dir,
        epochs=args.epochs,
        batch_size=args.batch_size,
        device=device,
        num_workers=args.num_workers,
        precision=args.precision,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        evaluation_interval=args.evaluation_interval,
        checkpoint_interval=args.checkpoint_interval,
        seed=args.seed,
        metric_logger=metric_logger,
    )
    if args.resume_from is None:
        model_trainer.train()
    else:
        model_trainer.resume(args.resume_from)


if __name__ == "__main__":
    train_yolox(parse_args())

Launch Training

Start a new run with the defaults:

python train_yolox.py \
  --dataset-dir path/to/dataset \
  --output-dir results/yolox

Override the most important model and Trainer settings as needed:

python train_yolox.py \
  --dataset-dir path/to/dataset \
  --output-dir results/yolox \
  --run-name custom-yolox \
  --variant medium \
  --epochs 300 \
  --batch-size 8 \
  --learning-rate 0.01 \
  --precision fp32

Resume Training

Resume to the configured total epoch count while keeping the model configuration compatible:

python train_yolox.py \
  --dataset-dir path/to/dataset \
  --output-dir results/yolox \
  --resume-from results/yolox/latest.pt

Command-Line Arguments

The training script accepts the following options:

Argument	Default	Description
--dataset-dir	required	COCO dataset root or split directory
--output-dir	required	Training outputs, logs, and checkpoints
--run-name	"yolox"	Name stored in run and export metadata
--variant	"small"	YOLOX architecture variant
--resolution	None	Square resolution override; otherwise uses the variant default
--pretrained / --no-pretrained	enabled	Enable or disable upstream pretrained weights
--epochs	300	Total training epochs, including restored epochs
--batch-size	8	Samples per micro-batch
--gradient-accumulation-steps	1	Micro-batches accumulated per optimizer update
--num-workers	2	DataLoader worker processes
--learning-rate	0.01	SGD learning rate
--weight-decay	5e-4	SGD weight decay
--nms-threshold	0.65	NMS IoU threshold used during evaluation
--max-detections	300	Maximum evaluation detections per image
--no-aug-epochs	15	Final epochs without multi-image augmentation
--warmup-epochs	5.0	Minimum scheduler warmup in epochs
--validation-split	0.3	Fallback validation fraction when no validation split exists
--evaluation-interval	10	Epochs between validation runs
--checkpoint-interval	10	Epochs between numbered checkpoints
--seed	42	Training and dataset-split seed
--device	automatic	Device such as cuda, cuda:0, or cpu
--precision	"fp32"	"auto", "fp32", "fp16", or "bf16"
--resume-from	None	Trainer checkpoint from which to continue

Training Guidance

See Best Practices: Trainer for guidance on batching, gradient accumulation, precision, evaluation, and checkpoints.

Run the Example

Run the maintained example from the Telekinesis examples repository:

cd telekinesis-examples
python examples/iris/train_yolox.py \
  --dataset-dir path/to/dataset \
  --output-dir results/yolox

Add --resume-from results/yolox/latest.pt to continue a previous run.

Category IDs

category_ids must contain exactly num_classes unique, non-negative COCO IDs. The adapter maps them to contiguous YOLOX indices during training and maps predictions back to the original IDs. Every dataset label must appear in this tuple.

Model Reference

Model Variants

Variant	Value	Default resolution	Depthwise
YOLOXVariant.NANO	"nano"	416	Yes
YOLOXVariant.TINY	"tiny"	416	No
YOLOXVariant.SMALL	"small"	640	No
YOLOXVariant.MEDIUM	"medium"	640	No
YOLOXVariant.LARGE	"large"	640	No
YOLOXVariant.XLARGE	"xlarge"	640	No

YOLOXConfig

YOLOXConfig combines architecture, category mapping, pretrained weights, augmentation, EMA, multiscale training, evaluation, and optimization settings.

from telekinesis.iris.models import YOLOX, YOLOXConfig, YOLOXVariant
from telekinesis.iris.optimizer import OptimizationConfig

config = YOLOXConfig(
    variant=YOLOXVariant.SMALL,
    num_classes=2,
    category_ids=(1, 2),
    pretrained=True,
    optimization=OptimizationConfig(
        optimizer="sgd",
        learning_rate=0.01,
        weight_decay=5e-4,
        optimizer_options={"momentum": 0.9, "nesterov": True},
        scheduler="cosine",
        scheduler_options={"min_factor": 0.05},
        scheduler_interval="step",
    ),
)
model = YOLOX(config)

Option	Default	Description
variant	YOLOXVariant.SMALL	Architecture size
num_classes	80	Number of predicted classes
category_ids	range(num_classes)	Ordered mapping from internal indices to COCO category IDs
pretrained	False	Loads upstream YOLOX weights; a replaced class head is allowed
checkpoint_path	None	Optional local pretrained checkpoint instead of the upstream URL
resolution	variant default	Square resolution override; at least 32 and divisible by 32
activation	"silu"	"silu", "relu", or "lrelu"
augmentation	YOLOX_DEFAULTS	Target-aware augmentation preset or declarative pipeline
nms_threshold	0.65	NMS threshold used by adapter evaluation
max_detections	300	Maximum evaluation detections
use_ema	True	Evaluates and exports an exponential moving average model
ema_decay	0.9998	EMA decay coefficient
no_aug_epochs	15	Final epochs without data composition; enables the L1 head loss
warmup_epochs	5.0	Minimum scheduler warmup expressed in epochs
multi_scale	True	Changes training resolution every ten optimizer updates
optimization	YOLOX SGD/cosine defaults	OptimizationConfig used by the adapter

The default optimization uses SGD at 0.01, weight decay 5e-4, momentum 0.9, Nesterov updates, a per-step cosine schedule, and minimum factor 0.05.

Augmentation and Training Behavior

YOLOX_DEFAULTS applies horizontal flip, photometric distortion, Mosaic, MixUp, and letterboxing. The adapter wraps the training dataset in MultiImageAugmentationDataset, disables composition for the last no_aug_epochs, enables the YOLOX L1 loss during that phase, and evaluates the EMA model when enabled.

Use a lower-resolution variant and smaller batch while validating a pipeline. Disable multi_scale for fixed-shape debugging, and set augmentation=() for a minimal pipeline.

Checkpoints and Export

YOLOX selects best.pt using validation mAP_50_95. Checkpoints also store EMA state, category IDs, run metadata, and a deployment state dict. Export the selected checkpoint to ONNX:

from telekinesis.iris import export

artifact = export.export_model(
    "results/yolox/best.pt",
    "artifacts/yolox",
    options=export.YOLOXONNXExportOptions(dynamic_batch=False),
)

Next Steps

Export YOLOX

Export the best checkpoint to ONNX with embedded metadata.

Export a model →

Run Inference

Run YOLOX through the unified deployment API.

Run inference →