Training, Inference, Evaluation, and Metrics API

Extension point. Task trainers, inferers, evaluators, and metric functions are implementation-level contracts between a task wrapper and its data/model route.

Runtime module map

Area	Modules and non-private entry points	Role
Classification	Trainer_Classification and trainer subclasses	Selects video, tracking, or frame training behavior.
Localization	build_trainer, build_inferer, build_evaluator, Trainer_*, Inferer, Evaluator	Routes model families and runners.
VQA	Trainer_VQA, LoRA trainer classes, SFT datasets/collators	Implements backend-specific fine-tuning and inference preparation.
Classification metrics	process_preds_labels, compute_classification_metrics	Computes classification scores from logits or labels.
Localization metrics	AP/mAP functions, prediction processors, result writers	Computes spotting metrics and may write evaluation artifacts.
VQA metrics	compute_vqa_metrics	Computes answer matching and configured referee-semantic metrics.

store_eval_files_json, detailed classification metrics, and trainer/checkpoint operations write files under their supplied/configured directories. Evaluation functions may require ground truth in the task’s expected representation.

Classification runtime

Trainer_Classification selects one concrete trainer based on the configured input route. BaseTrainerClassification provides common behavior; MVTrainerClassification, TrackingTrainerClassification, and FramesTrainerClassification implement the multi-view, tracking, and frame paths respectively. These classes are internal; ClassificationModel is the stable operation entry point.

classification trainers for video and tracking modalities.

provides a base trainer with modality-agnostic training, validation, and test loops, plus two modality-specific subclasses that implement the forward pass. Trainer_Classification is the top-level dispatcher consumed by the API layer.

BaseTrainerClassification

modality-agnostic training loop for classification.

handles epoch iteration, gradient updates, DDP gather, metric computation, W&B logging, checkpoint saving, and JSON prediction export. subclasses only need to override _forward_batch() with modality-specific tensor preparation.

Parameters:

Name	Type	Description	Default
train_loader		DataLoader for training set.	required
val_loader		DataLoader for validation set.	required
test_loader		DataLoader for test set (may be None during training).	required
model		the classification model (already on device).	required
optimizer		PyTorch optimizer.	required
scheduler		learning-rate scheduler.	required
criterion		loss function callable.	required
class_weights		optional per-class weight tensor for the loss.	required
class_names		dict mapping class indices to names.	required
save_dir		root directory for checkpoint and prediction output.	required
model_name		name used for the checkpoint sub-directory.	required
max_epochs		maximum number of training epochs.	1000
device		torch.device or device string.	'cuda'
top_k		k value for top-k accuracy computation.	2
wandb_project		W&B project name.	required
wandb_run_name		W&B run display name.	required
wandb_config		dict of hyperparameters logged to W&B.	required
patience		early-stopping patience (0=disabled).	10
monitor		metric name to monitor for checkpointing.	'balanced_accuracy'
mode		"max" or "min" depending on the monitored metric.	'max'

Source code in opensportslib/core/trainer/classification_trainer.py

class BaseTrainerClassification:
    """modality-agnostic training loop for classification.

    handles epoch iteration, gradient updates, DDP gather, 
    metric computation, W&B logging, checkpoint saving, and JSON
    prediction export. subclasses only need to override _forward_batch()
    with modality-specific tensor preparation.

    Args:
        train_loader: DataLoader for training set.
        val_loader: DataLoader for validation set.
        test_loader: DataLoader for test set (may be None during training).
        model: the classification model (already on device).
        optimizer: PyTorch optimizer.
        scheduler: learning-rate scheduler.
        criterion: loss function callable.
        class_weights: optional per-class weight tensor for the loss.
        class_names: dict mapping class indices to names.
        save_dir: root directory for checkpoint and prediction output.
        model_name: name used for the checkpoint sub-directory.
        max_epochs: maximum number of training epochs.
        device: torch.device or device string.
        top_k: k value for top-k accuracy computation.
        wandb_project: W&B project name.
        wandb_run_name: W&B run display name.
        wandb_config: dict of hyperparameters logged to W&B.
        patience: early-stopping patience (0=disabled).
        monitor: metric name to monitor for checkpointing.
        mode: "max" or "min" depending on the monitored metric.
    """

    def __init__(
        self,
        train_loader,
        val_loader,
        test_loader,
        model,
        optimizer,
        scheduler,
        criterion,
        class_weights,
        class_names,
        save_dir,
        model_name,
        max_epochs=1000,
        device="cuda",
        top_k=2,
        patience=10,
        monitor="balanced_accuracy",
        mode="max",
        revert_on_lr_reduction=False,
        config=None,
    ):
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.test_loader = test_loader

        self.model = model#.to(device)
        #self.model = DDP(self.model, device_ids=[device])
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.criterion = criterion
        self.class_weights = class_weights
        self.class_names = class_names

        self.model_name = model_name
        self.max_epochs = max_epochs
        self.device = device
        self.top_k = top_k
        self.patience = patience

        self.monitor = monitor
        self.mode = mode
        self.config = config

        self.best_checkpoint_path = None
        self.best_metric = None
        self.revert_on_lr_reduction = revert_on_lr_reduction
        self._best_model_state = None
        self.predictions_payload = None
        self.predictions_path = None

        self.rank = dist.get_rank() if dist.is_initialized() else 0

        self.save_dir = save_dir
        os.makedirs(self.save_dir, exist_ok=True)

        try:
            if self.rank == 0:
                wandb.watch(self.model, log="gradients", log_freq=100)
        except Exception:
            pass

    # -- abstract forward pass --------------------------------------

    def _forward_batch(self, batch):
        """run the modality-specific forward pass.

        must be overridden by every subclass.

        Args:
            batch: a dict produced by the DataLoader.

        Returns:
            a tuple (logits, labels) where both are tensors on
            self.device.
        """
        raise NotImplementedError

    # -- process batch ----------------------------------------------

    def _process_batch(self, batch, train):
        """run forward pass, compute loss, and optionally update weights.

        the default implementation calls _forward_batch() for the
        modality-specific forward pass, then computes the loss and
        runs the backward step.  subclasses may override this entirely
        to inject AMP, mixup, or other training-time modifications
        without touching the base training loop.

        Args:
            batch: a dict produced by the DataLoader.
            train: if True, compute gradients and update weights.

        Returns:
            a tuple (logits, labels, loss).
        """
        has_labels = "labels" in batch or "label" in batch
        with torch.set_grad_enabled(train):
            logits, labels = self._forward_batch(batch)
            if labels is None:
                    has_labels = False
            loss = None
            if has_labels:
                if self.class_weights is not None:
                    loss = self.criterion(
                        output=logits, labels=labels,
                        weight=self.class_weights.to(self.device)
                    )
                else:
                    loss = self.criterion(output=logits, labels=labels)

            if train and loss is not None:
                self.optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.optimizer.step()

        return logits, labels, loss, has_labels

    # -- training loop ----------------------------------------------

    def train(self, epoch_start=0, save_every=3):
        """run the full training loop with validation after each epoch.

        Args:
            epoch_start: the epoch number to start from (0-based).
            save_every: currently unused; reserved for periodic 
            checkpoint saving.
        """
        logging.info("Starting training")
        monitor = self.monitor
        mode = self.mode
        best_metric = -float("inf") if mode == "max" else float("inf")
        best_path = None

        for epoch in range(epoch_start, self.max_epochs):
            logging.info(f"\nEpoch {epoch+1}/{self.max_epochs}")

            # --- train ---
            if hasattr(self.train_loader.sampler, "set_epoch"):
                self.train_loader.sampler.set_epoch(epoch)

            disable = self.rank != 0

            pbar = tqdm.tqdm(
                total=len(self.train_loader), desc="Training", 
                position=0, leave=True, disable=disable
            )
            _, _, train_loss, train_metrics = self._run_epoch(
                self.train_loader, epoch + 1, 
                train=True, set_name="train", pbar=pbar
            )
            pbar.close()

            # --- validation ---
            pbar = tqdm.tqdm(
                total=len(self.val_loader), desc="Valid", 
                position=1, leave=True, disable=disable
            )
            _, _, val_loss, val_metrics = self._run_epoch(
                self.val_loader, epoch + 1, 
                train=False, set_name="valid", pbar=pbar
            )
            pbar.close()

            prev_lr = self.optimizer.param_groups[0]["lr"]

            # capture LR before the scheduler step so we can detect
            # plateau-triggered reductions.
            val_metric = val_metrics.get(
                "balanced_accuracy", val_metrics.get("accuracy", 0)
            )
            train_metric = train_metrics.get(
                "balanced_accuracy", train_metrics.get("accuracy", 0)
            )

            # ReduceLROnPlateau needs the monitored metric
            if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                self.scheduler.step(val_loss)
            else:
                self.scheduler.step()

            current_lr = self.optimizer.param_groups[0]["lr"]

            # early stopping: mirror pixels_vs_positions behavior
            if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                min_lr = self.scheduler.min_lrs[0]
                if current_lr <= 2 * min_lr:
                    if self.rank == 0:
                        logging.info(
                            f"Early stopping at epoch {epoch+1}: "
                            f"lr {current_lr:.2e} <= 2 * min_lr {min_lr:.2e}"
                        )
                    break

            # When ReduceLROnPlateau drops the LR, revert weights to
            # the best checkpoint so training continues from the
            # strongest point rather than from a potentially overfit
            # state. This mirrors the pixels_vs_positions recipe.
            if (
                self.revert_on_lr_reduction
                and current_lr != prev_lr
                and self._best_model_state is not None
            ):
                self.model.load_state_dict(self._best_model_state)
                print(
                    f"LR reduced from {prev_lr:.2e} to {current_lr:.2e} "
                    f"-- reverted to best model"
                )

            if self.rank == 0:
                # ---------------- W&B LOG ----------------
                if wandb.run is not None:
                    wandb.log({
                        "epoch": epoch + 1,
                        "lr": current_lr,
                        "train/loss": train_loss,
                        "valid/loss": val_loss,
                        **{f"train/{k}": v for k, v in train_metrics.items()},
                        **{f"valid/{k}": v for k, v in val_metrics.items()},
                    })

                logging.info(f"Train Loss: {train_loss:.4f} | Train Bal Acc: {train_metric:.4f}")
                logging.info(f"Val Loss: {val_loss:.4f} | Val Bal Acc: {val_metric:.4f}")

            # ---------------- CHECKPOINT ----------------
            current = val_loss if monitor == "loss" else val_metrics.get(monitor, 0)
            if dist.is_initialized():
                current_tensor = torch.tensor(
                    [float(current if self.rank == 0 else 0.0)],
                    device=self.device,
                )
                dist.broadcast(current_tensor, src=0)
                current = current_tensor.item()

            is_better = current > best_metric if mode == "max" else current < best_metric

            if is_better:
                best_metric = current
                self.best_metric = best_metric

                if self.revert_on_lr_reduction:
                    self._best_model_state = {
                        k: v.cpu().clone()
                        for k, v in self.model.state_dict().items()
                    }

                if self.rank == 0:
                    best_path = self._save_checkpoint("best", epoch + 1, tag="best")
                    self.best_checkpoint_path = best_path

                    if wandb.run is not None:
                        artifact = wandb.Artifact("model-checkpoint", type="model")
                        artifact.add_file(best_path)
                        wandb.log_artifact(artifact)

        if self.rank == 0:
            logging.info(f"Best checkpoint : {self.best_checkpoint_path}")
            logging.info("Training finished.")


    # -- TEST evaluation ------------------------------------------

    def test(self, epoch=None, detailed_results=False):
        """run the test set evaluation.

        Args:
            epoch: the epoch number to evaluate (if None, uses "final").
            detailed_results: whether to compute detailed classification metrics.

        Returns:
            a tuple (test_loss, test_metrics).
        """
        logging.info("\nRunning TEST evaluation")
        pbar = tqdm.tqdm(
            total=len(self.test_loader), desc="Test", position=0, 
            leave=True, disable = self.rank != 0
        )
        all_logits, all_labels, test_loss, test_metrics = self._run_epoch(
            self.test_loader,
            epoch if epoch is not None else "final",
            train=False, set_name="test", pbar=pbar
        )
        pbar.close()

        if self.rank==0:
            if wandb.run is not None:
                wandb.log({
                    "test/loss": test_loss,
                    **{f"test/{k}": v for k, v in test_metrics.items()},
                })

            if detailed_results:
                from opensportslib.metrics.classification_metric import (
                    compute_detailed_classification_metrics
                )
                compute_detailed_classification_metrics(
                    all_logits=all_logits, all_labels=all_labels,
                    class_names=self.class_names, save_dir=self.save_dir,
                    set_name="test"
                )

        logging.info(f"TEST METRICS : {test_metrics}")
        return test_loss, test_metrics

    # -- single epoch logic -------------------------------------

    def _run_epoch(self, dataloader, epoch, train=False, set_name="train", pbar=None):
        """execute one pass over a dataloader.

        handles forward/backward, per-batch bookkeeping, DDP gather, metric
        computation, confusion-matrix logging, and JSON prediction export.

        Args:
            dataloader: the DataLoader to iterate over.
            epoch: the epoch number (for checkpointing and folder naming).
            train: if True, compute gradients and update weights.
            set_name: "train", "valid", or "test" (for logging and JSON).
            pbar: optional tqdm progress bar.

        Returns:
            a tuple (all_logits, all_labels, avg_loss, metrics).
            on non-rank-0 DDP workers the first two are None and metrics
            is an empty dict.
        """

        import torch.distributed as dist

        if train:
            self.model.train()
        else:
            self.model.eval()

        total_loss = 0.0
        total_batches = 0

        all_logits = []
        all_labels = []
        results = []

        # -------- Create epoch folder --------
        epoch_dir = os.path.join(self.save_dir, str(epoch))
        os.makedirs(epoch_dir, exist_ok=True)
        save_path = os.path.join(
            epoch_dir, f"predictions_{set_name}_epoch_{epoch}.json"
        )

        # --- batch loop ---
        for batch in dataloader:
            if pbar:
                pbar.update()

            logits, labels, loss, has_labels = self._process_batch(batch, train)

            if loss is not None:
                total_loss += loss.item()
                total_batches += 1

            logits_cpu = logits.detach().cpu()
            all_logits.append(logits_cpu)

            if has_labels:
                labels_cpu = labels.detach().cpu()
                all_labels.append(labels_cpu)

            # per-sample predictions for JSON export.
            probs = torch.softmax(logits_cpu, dim=1)
            preds = torch.argmax(probs, dim=1)
            confs = probs.max(dim=1).values
            ids = batch["id"]

            for i in range(len(preds)):
                results.append({
                    "id": ids[i],
                    "pred_label": self.class_names[preds[i].item()],
                    "confidence": float(confs[i].item()),
                    "pred_class_idx": preds[i].item(),
                })

        # --- concatenate local predictions ---
        if len(all_logits) > 0:
            all_logits = torch.cat(all_logits).numpy()
        else:
            all_logits = np.zeros((0, 1))

        if len(all_labels) > 0:
            all_labels = torch.cat(all_labels).numpy()
        else:
            all_labels = np.zeros((0,))

        avg_loss = total_loss / max(1, total_batches)

        # --- DDP gather (handles uneven shard sizes) ---
        if dist.is_initialized():
            loss_stats = torch.tensor(
                [total_loss, total_batches],
                dtype=torch.float64,
                device=self.device,
            )
            dist.all_reduce(loss_stats, op=dist.ReduceOp.SUM)
            avg_loss = (
                loss_stats[0] / torch.clamp(loss_stats[1], min=1.0)
            ).item()

            gathered = [None for _ in range(dist.get_world_size())]
            dist.all_gather_object(gathered, (all_logits, all_labels, results))

            if self.rank == 0:
                all_logits = np.concatenate([g[0] for g in gathered])
                all_labels = np.concatenate([g[1] for g in gathered])
                results = [r for g in gathered for r in g[2]]
            else:
                self.predictions_payload = None
                return None, None, avg_loss, {}

        # --- metrics (rank-0 only in DDP) ---
        if len(all_labels) > 0:
            metrics = compute_classification_metrics(
                (all_logits, all_labels), top_k=self.top_k,
            )
        else:
            metrics = {}

        # --- confusion matrix (validation and test only) ---
        if self.rank == 0 and set_name in ["valid", "test"] and len(all_labels) > 0:
            preds_all, labels_all, _ = process_preds_labels(
                (all_logits, all_labels)
            )
            class_names = [
                self.class_names[i] for i in sorted(self.class_names.keys())
            ]

            log_confusion_matrix_wandb(
                y_true=labels_all.tolist(),
                y_pred=preds_all.tolist(),
                class_names=class_names,
                split_name=set_name,
            )

        # --- save JSON (rank-0 only) ---
        if self.rank == 0:
            submission = {
                "version": "2.0",
                "task": "action_classification",
                "date": datetime.now().strftime("%Y-%m-%d"),
                "metadata": {"type": "predictions"},
                "data": [],
            }

            for r in results:
                submission["data"].append({
                    "id": r["id"],
                    "labels": {
                        "action": {
                            "label": r["pred_label"],
                            "confidence": r["confidence"],
                        }
                    },
                })

            logging.info(f"RESULTS Length: {len(results)}")
            logging.info(f"Predictions are stored at : {save_path}")
            with open(save_path, "w", encoding="utf-8") as f:
                json.dump(submission, f, indent=2)
            self.predictions_payload = submission
            self.predictions_path = save_path

        return all_logits, all_labels, avg_loss, metrics


    # -- checkpoint saving ---------------------------------------

    def _save_checkpoint(self, filename, epoch, tag=None):
        epoch_dir = os.path.join(self.save_dir, str(filename))
        os.makedirs(epoch_dir, exist_ok=True)

        state = {
            "epoch": epoch,
            "state_dict": self.model.module.state_dict() if hasattr(self.model, 'module') else self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "scheduler": self.scheduler.state_dict(),
            "monitor": self.monitor,
            "mode": self.mode,
            "best_metric": self.best_metric,
        }

        if hasattr(self, "scaler"):
            state["scaler"] = self.scaler.state_dict()

        name = f"epoch_{epoch}.pt"
        if tag:
            name = f"{tag}_epoch_{epoch}.pt"

        path_aux = os.path.join(epoch_dir, name)
        torch.save(state, path_aux)
        logging.info(f"Saved checkpoint: {path_aux}")

        if self.config is not None:
            from opensportslib.core.utils.config import save_config
            config_path = os.path.join(epoch_dir, "config.yaml")
            save_config(self.config, config_path)
            logging.info(f"Saved config: {config_path}")

        return path_aux

test(epoch=None, detailed_results=False)

run the test set evaluation.

Parameters:

Name	Type	Description	Default
epoch		the epoch number to evaluate (if None, uses "final").	None
detailed_results		whether to compute detailed classification metrics.	False

Returns:

Type	Description
	a tuple (test_loss, test_metrics).

Source code in opensportslib/core/trainer/classification_trainer.py

def test(self, epoch=None, detailed_results=False):
    """run the test set evaluation.

    Args:
        epoch: the epoch number to evaluate (if None, uses "final").
        detailed_results: whether to compute detailed classification metrics.

    Returns:
        a tuple (test_loss, test_metrics).
    """
    logging.info("\nRunning TEST evaluation")
    pbar = tqdm.tqdm(
        total=len(self.test_loader), desc="Test", position=0, 
        leave=True, disable = self.rank != 0
    )
    all_logits, all_labels, test_loss, test_metrics = self._run_epoch(
        self.test_loader,
        epoch if epoch is not None else "final",
        train=False, set_name="test", pbar=pbar
    )
    pbar.close()

    if self.rank==0:
        if wandb.run is not None:
            wandb.log({
                "test/loss": test_loss,
                **{f"test/{k}": v for k, v in test_metrics.items()},
            })

        if detailed_results:
            from opensportslib.metrics.classification_metric import (
                compute_detailed_classification_metrics
            )
            compute_detailed_classification_metrics(
                all_logits=all_logits, all_labels=all_labels,
                class_names=self.class_names, save_dir=self.save_dir,
                set_name="test"
            )

    logging.info(f"TEST METRICS : {test_metrics}")
    return test_loss, test_metrics

train(epoch_start=0, save_every=3)

run the full training loop with validation after each epoch.

Parameters:

Name	Type	Description	Default
epoch_start		the epoch number to start from (0-based).	0
save_every		currently unused; reserved for periodic	3

Source code in opensportslib/core/trainer/classification_trainer.py

def train(self, epoch_start=0, save_every=3):
    """run the full training loop with validation after each epoch.

    Args:
        epoch_start: the epoch number to start from (0-based).
        save_every: currently unused; reserved for periodic 
        checkpoint saving.
    """
    logging.info("Starting training")
    monitor = self.monitor
    mode = self.mode
    best_metric = -float("inf") if mode == "max" else float("inf")
    best_path = None

    for epoch in range(epoch_start, self.max_epochs):
        logging.info(f"\nEpoch {epoch+1}/{self.max_epochs}")

        # --- train ---
        if hasattr(self.train_loader.sampler, "set_epoch"):
            self.train_loader.sampler.set_epoch(epoch)

        disable = self.rank != 0

        pbar = tqdm.tqdm(
            total=len(self.train_loader), desc="Training", 
            position=0, leave=True, disable=disable
        )
        _, _, train_loss, train_metrics = self._run_epoch(
            self.train_loader, epoch + 1, 
            train=True, set_name="train", pbar=pbar
        )
        pbar.close()

        # --- validation ---
        pbar = tqdm.tqdm(
            total=len(self.val_loader), desc="Valid", 
            position=1, leave=True, disable=disable
        )
        _, _, val_loss, val_metrics = self._run_epoch(
            self.val_loader, epoch + 1, 
            train=False, set_name="valid", pbar=pbar
        )
        pbar.close()

        prev_lr = self.optimizer.param_groups[0]["lr"]

        # capture LR before the scheduler step so we can detect
        # plateau-triggered reductions.
        val_metric = val_metrics.get(
            "balanced_accuracy", val_metrics.get("accuracy", 0)
        )
        train_metric = train_metrics.get(
            "balanced_accuracy", train_metrics.get("accuracy", 0)
        )

        # ReduceLROnPlateau needs the monitored metric
        if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
            self.scheduler.step(val_loss)
        else:
            self.scheduler.step()

        current_lr = self.optimizer.param_groups[0]["lr"]

        # early stopping: mirror pixels_vs_positions behavior
        if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
            min_lr = self.scheduler.min_lrs[0]
            if current_lr <= 2 * min_lr:
                if self.rank == 0:
                    logging.info(
                        f"Early stopping at epoch {epoch+1}: "
                        f"lr {current_lr:.2e} <= 2 * min_lr {min_lr:.2e}"
                    )
                break

        # When ReduceLROnPlateau drops the LR, revert weights to
        # the best checkpoint so training continues from the
        # strongest point rather than from a potentially overfit
        # state. This mirrors the pixels_vs_positions recipe.
        if (
            self.revert_on_lr_reduction
            and current_lr != prev_lr
            and self._best_model_state is not None
        ):
            self.model.load_state_dict(self._best_model_state)
            print(
                f"LR reduced from {prev_lr:.2e} to {current_lr:.2e} "
                f"-- reverted to best model"
            )

        if self.rank == 0:
            # ---------------- W&B LOG ----------------
            if wandb.run is not None:
                wandb.log({
                    "epoch": epoch + 1,
                    "lr": current_lr,
                    "train/loss": train_loss,
                    "valid/loss": val_loss,
                    **{f"train/{k}": v for k, v in train_metrics.items()},
                    **{f"valid/{k}": v for k, v in val_metrics.items()},
                })

            logging.info(f"Train Loss: {train_loss:.4f} | Train Bal Acc: {train_metric:.4f}")
            logging.info(f"Val Loss: {val_loss:.4f} | Val Bal Acc: {val_metric:.4f}")

        # ---------------- CHECKPOINT ----------------
        current = val_loss if monitor == "loss" else val_metrics.get(monitor, 0)
        if dist.is_initialized():
            current_tensor = torch.tensor(
                [float(current if self.rank == 0 else 0.0)],
                device=self.device,
            )
            dist.broadcast(current_tensor, src=0)
            current = current_tensor.item()

        is_better = current > best_metric if mode == "max" else current < best_metric

        if is_better:
            best_metric = current
            self.best_metric = best_metric

            if self.revert_on_lr_reduction:
                self._best_model_state = {
                    k: v.cpu().clone()
                    for k, v in self.model.state_dict().items()
                }

            if self.rank == 0:
                best_path = self._save_checkpoint("best", epoch + 1, tag="best")
                self.best_checkpoint_path = best_path

                if wandb.run is not None:
                    artifact = wandb.Artifact("model-checkpoint", type="model")
                    artifact.add_file(best_path)
                    wandb.log_artifact(artifact)

    if self.rank == 0:
        logging.info(f"Best checkpoint : {self.best_checkpoint_path}")
        logging.info("Training finished.")

MVTrainerClassification

Bases: BaseTrainerClassification

forward pass for multi-view video classification.

expects batches with pixel_values of shape (B, V, C, T, H, W) and integer labels of shape (B,).

Source code in opensportslib/core/trainer/classification_trainer.py

class MVTrainerClassification(BaseTrainerClassification):
    """forward pass for multi-view video classification.

    expects batches with pixel_values of shape
    (B, V, C, T, H, W) and integer labels of shape (B,).
    """

    def _forward_batch(self, batch):
        """move video clips to device and run the model.

        Args:
            batch: dict with keys "pixel_values" and "labels".

        Returns:
            a tuple (logits, labels) on self.device.
        """
        mvclips = batch["pixel_values"].to(self.device).float()
        labels = batch.get("labels", None)
        if labels is not None:
            labels = labels.to(self.device)

        outputs = self.model(mvclips)

        if isinstance(outputs, tuple):
            logits = outputs[0]
        else:
            logits = outputs

        if logits.dim() == 1:
            logits = logits.unsqueeze(0)

        return logits, labels

TrackingTrainerClassification

Bases: BaseTrainerClassification

forward pass for tracking-based classification.

expects batches with x of shape (B, N, 2), edge_index of shape (2, E), batch of shape (B,), batch_size, seq_len, and integer labels of shape (B,).

Source code in opensportslib/core/trainer/classification_trainer.py

class TrackingTrainerClassification(BaseTrainerClassification):
    """forward pass for tracking-based classification.

    expects batches with x of shape (B, N, 2), edge_index of shape (2, E),
    batch of shape (B,), batch_size, seq_len, and integer labels of shape (B,).
    """

    def _forward_batch(self, batch):
        """move tracking data to device and run the model.

        Args:
            batch: dict with keys "x", "edge_index", "batch", "batch_size",
            "seq_len", and "labels".

        Returns:
            a tuple (logits, labels) on self.device.
        """
        tracking_batch = {
            "x": batch["x"].to(self.device),
            "edge_index": batch["edge_index"].to(self.device),
            "batch": batch["batch"].to(self.device),
            "batch_size": batch["batch_size"],
            "seq_len": batch["seq_len"],
        }
        labels = batch.get("labels", None)
        if labels is not None:
            labels = labels.to(self.device)

        logits = self.model(tracking_batch)

        return logits, labels

FramesTrainerClassification

Bases: BaseTrainerClassification

forward pass for frames_npy video classification.

supports optional mixed-precision training (AMP) and mixup augmentation, controlled via config.TRAIN.use_amp and config.TRAIN.mixup_alpha respectively.

expects batches with pixel_values of shape (B, T, H, W, C) and integer labels of shape (B,).

Source code in opensportslib/core/trainer/classification_trainer.py

class FramesTrainerClassification(BaseTrainerClassification):
    """forward pass for frames_npy video classification.

    supports optional mixed-precision training (AMP) and mixup
    augmentation, controlled via config.TRAIN.use_amp and
    config.TRAIN.mixup_alpha respectively.

    expects batches with pixel_values of shape (B, T, H, W, C)
    and integer labels of shape (B,).
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        cfg = self.config
        self.use_amp = getattr(cfg.TRAIN, "use_amp", False) if cfg else False
        self.mixup_alpha = getattr(cfg.TRAIN, "mixup_alpha", 0.0) if cfg else 0.0
        self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_amp)

    def _forward_batch(self, batch):
        pixel_values = batch["pixel_values"].to(self.device).float()
        labels = batch["labels"].to(self.device)
        logits = self.model({"pixel_values": pixel_values})
        return logits, labels

    def _process_batch(self, batch, train):
        pixel_values = batch["pixel_values"].to(self.device).float()
        labels = batch["labels"].to(self.device)

        with torch.set_grad_enabled(train):
            use_mixup = (
                train
                and self.mixup_alpha > 0
                and np.random.random() > 0.5
            )

            with torch.amp.autocast("cuda", enabled=self.use_amp):
                if use_mixup:
                    pixel_values, labels_a, labels_b, lam = mixup_data(
                        pixel_values, labels, self.mixup_alpha
                    )
                    logits = self.model({"pixel_values": pixel_values})
                    loss = (
                        lam * self.criterion(output=logits, labels=labels_a)
                        + (1 - lam) * self.criterion(output=logits, labels=labels_b)
                    )
                    labels = labels_a
                else:
                    logits = self.model({"pixel_values": pixel_values})
                    if self.class_weights is not None:
                        loss = self.criterion(
                            output=logits, labels=labels,
                            weight=self.class_weights.to(self.device),
                        )
                    else:
                        loss = self.criterion(output=logits, labels=labels)

            if train:
                self.optimizer.zero_grad(set_to_none=True)
                self.scaler.scale(loss).backward()
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.scaler.step(self.optimizer)
                self.scaler.update()

        return logits, labels, loss, True

Trainer_Classification

high-level trainer that dispatches to the right modality trainer.

consumed by ClassificationModel. Responsible for building data loaders, optimizers, schedulers, and samplers, then delegating the actual loop to MVTrainerClassification or TrackingTrainerClassification.

Parameters:

Name	Type	Description	Default
config		the configuration object.	required

Source code in opensportslib/core/trainer/classification_trainer.py

class Trainer_Classification:
    """high-level trainer that dispatches to the right modality trainer.

    consumed by ClassificationModel. Responsible for building data
    loaders, optimizers, schedulers, and samplers, then delegating the
    actual loop to MVTrainerClassification or TrackingTrainerClassification.

    Args:
        config: the configuration object.
    """

    def __init__(self, config):
        self.config = config
        self.device = select_device(self.config.SYSTEM)
        self.save_dir = get_system_path(config, "save_dir", "./checkpoints") or "./checkpoints"
        self.model = None
        self.optimizer = None
        self.scheduler = None
        self.epoch = 0
        self.trainer = None
        self.predictions_payload = None

    def compute_metrics(self, pred, mode="logits"):
        """thin wrapper around the metric module.

        Args:
            pred: a tuple (logits, labels).
            mode: "logits" or "labels" (default: "logits").

        Returns:
            a dictionary of classification metrics.
        """
        return compute_classification_metrics(
            pred, top_k=2, mode=mode
        )

    # -- training -----------------------------------------------

    def train(self, model, train_dataset, val_dataset=None, rank=0, world_size=1):
        """build all training components and run the loop.

        detects the model type (HuggingFace vs. custom) and the data 
        modality (video vs. tracking) to select the right trainer class,
        sampler, and collate function.

        Args:
            model: the classification model.
            train_dataset: training ClassificationDataset.
            val_dataset: validation ClassificationDataset (optional).
            rank: GPU rank (0-indexed).
            world_size: total number of GPUs.
        """
        from opensportslib.core.loss.builder import build_criterion
        from opensportslib.core.optimizer.builder import build_optimizer
        from opensportslib.core.scheduler.builder import build_scheduler
        from opensportslib.core.utils.data import tracking_collate_fn
        from torch.nn.parallel import DistributedDataParallel as DDP
        from torch.utils.data.distributed import DistributedSampler

        is_ddp = world_size > 1
        modality = get_data_modality(self.config)
        is_tracking_modality = _is_tracking_graph_modality(modality)
        is_frames_modality = _is_frames_npy_modality(self.config, modality)
        seed = get_system_seed(self.config)

        g = torch.Generator()
        g.manual_seed(seed)

        # HuggingFace-backed models use the HF Trainer path when routed there.

        if is_ddp:
                torch.cuda.set_device(rank)
                self.device = torch.device(f"cuda:{rank}")
        else:
                self.device = select_device(self.config.SYSTEM)

        self.model = model.to(self.device)

        if is_ddp:
            self.model = DDP(self.model, device_ids=[rank])

        # Build components
        optimizer = build_optimizer(
            self.model.parameters(), cfg=self.config.TRAIN.optimizer
        )
        scheduler = build_scheduler(
            optimizer, cfg=self.config.TRAIN.scheduler
        )
        criterion = build_criterion(self.config.TRAIN.criterion)
        train_sampling = get_train_sampling(self.config)
        train_selection = get_train_selection(self.config)
        train_checkpoint = get_train_checkpoint(self.config)

        # --- class weights for the loss ---
        if train_sampling.get("use_weighted_loss", False):
            class_weights = train_dataset.get_class_weights(
                num_classes=train_dataset.num_classes(), sqrt=True
            ).to(self.device)
        else:
            class_weights = None

        # tracking modality needs a customm collate that merges PyG
        # Data objects into a single batched graph per timestamp.
        collate_fn = tracking_collate_fn if is_tracking_modality else None

        # --- train sampler ---
        if train_sampling.get("use_weighted_sampler", False):
            sample_weights = train_dataset.get_sample_weights()

            samples_per_class = train_sampling.get("samples_per_class")
            if samples_per_class:
                num_classes = train_dataset.num_classes()
                num_samples = samples_per_class * num_classes
            else:
                num_samples = len(sample_weights)

            if is_ddp:
                train_sampler = DistributedWeightedSampler(
                    weights=sample_weights,
                    num_replicas=world_size,
                    rank=rank,
                    replacement=True,
                    num_samples=num_samples,
                    seed=seed
                )
            else:
                train_sampler = WeightedRandomSampler(
                    weights=sample_weights,
                    num_samples=num_samples,
                    replacement=True,
                    generator=g
                )

            shuffle = False

        else:
            if is_ddp:
                train_sampler = DistributedSampler(
                    train_dataset,
                    num_replicas=world_size,
                    rank=rank,
                    shuffle=True,
                    drop_last=True
                )
            else:
                train_sampler = None

            shuffle = not is_ddp


        # --- validation sampler ---
        if is_ddp:
            val_sampler = DistributedSampler(
                val_dataset,
                num_replicas=world_size,
                rank=rank,
                shuffle=False,
                drop_last=False
            )
        else:
            val_sampler = None

        train_dataloader_cfg = get_split_dataloader_cfg(self.config, "train")
        train_num_workers = getattr(train_dataloader_cfg, "num_workers", 0)
        train_pin_memory = getattr(train_dataloader_cfg, "pin_memory", self.device.type == "cuda")
        train_mp_context = getattr(train_dataloader_cfg, "mp_context", None)
        train_persistent_workers = getattr(train_dataloader_cfg, "persistent_workers", train_num_workers > 0)
        train_prefetch_factor = getattr(train_dataloader_cfg, "prefetch_factor", 4 if train_num_workers > 0 else None)

        if train_mp_context is not None:
            train_mp_context = mp.get_context(train_mp_context)

        train_loader = DataLoader(
            train_dataset,
            batch_size=train_dataloader_cfg.batch_size,
            shuffle=(train_sampler is None and shuffle),
            sampler=train_sampler,
            num_workers=train_num_workers,
            pin_memory=train_pin_memory,
            collate_fn=collate_fn,
            worker_init_fn=seed_worker,
            generator=g,
            drop_last=True,
            multiprocessing_context=train_mp_context,
            persistent_workers=train_persistent_workers,
            prefetch_factor=train_prefetch_factor,
        )

        valid_dataloader_cfg = get_split_dataloader_cfg(self.config, "valid")
        valid_num_workers = getattr(valid_dataloader_cfg, "num_workers", 0)
        valid_pin_memory = getattr(valid_dataloader_cfg, "pin_memory", self.device.type == "cuda")
        valid_mp_context = getattr(valid_dataloader_cfg, "mp_context", None)
        valid_persistent_workers = getattr(valid_dataloader_cfg, "persistent_workers", valid_num_workers > 0)
        valid_prefetch_factor = getattr(valid_dataloader_cfg, "prefetch_factor", 4 if valid_num_workers > 0 else None)

        if valid_mp_context is not None:
            valid_mp_context = mp.get_context(valid_mp_context)

        val_loader = DataLoader(
            val_dataset,
            batch_size=valid_dataloader_cfg.batch_size,
            shuffle=False,
            sampler=val_sampler,   
            num_workers=valid_num_workers,
            pin_memory=valid_pin_memory,
            collate_fn=collate_fn,
            worker_init_fn=seed_worker,
            generator=g,
            multiprocessing_context=valid_mp_context,
            persistent_workers=valid_persistent_workers,
            prefetch_factor=valid_prefetch_factor,
        )

        # select the modality-specific trainer.
        if is_tracking_modality:
            TrainerClass = TrackingTrainerClassification
        elif is_frames_modality:
            TrainerClass = FramesTrainerClassification
        else:
            TrainerClass = MVTrainerClassification

        self.trainer = TrainerClass(
            train_loader=train_loader,
            val_loader=val_loader,
            test_loader=None,
            model=self.model,
            optimizer=optimizer,
            scheduler=scheduler,
            criterion=criterion,
            class_weights=class_weights,
            class_names=train_dataset.label_map,
            save_dir=self.save_dir,
            model_name=get_component_name_by_kind(self.config, "encoder") or "model",
            max_epochs=self.config.TRAIN.epochs,
            device=self.device,
            top_k=2,
            patience=train_selection.get("patience", 0),
            monitor=train_selection.get("monitor", "balanced_accuracy"),
            mode=train_selection.get("mode", "max"),
            revert_on_lr_reduction=(is_tracking_modality or is_frames_modality),
            config=self.config,
        )

        self.trainer.train(
            epoch_start=self.epoch,
            save_every=train_checkpoint.get("save_every", 1),
        )
        return getattr(self.trainer, "best_checkpoint_path", None)

    def _train_huggingface(self, model, train_dataset, val_dataset):
        """Handle HuggingFace Trainer for VideoMAE."""
        from opensportslib.core.sampler.weighted_sampler import WeightedTrainer, VideoMAETrainer

        self.model = model

        args = TrainingArguments(
            label_names=["labels"],
            output_dir=self.save_dir,
            per_device_train_batch_size=get_split_dataloader_cfg(self.config, "train").batch_size,
            per_device_eval_batch_size=get_split_dataloader_cfg(self.config, "valid").batch_size,
            num_train_epochs=self.config.TRAIN.epochs,
            eval_strategy="epoch" if val_dataset else "no",
            save_strategy="epoch",
            logging_strategy="steps",
            logging_steps=5,
            save_total_limit=10,
            load_best_model_at_end=True,
            fp16=True,
            warmup_ratio=0.1,
        )

        if get_train_sampling(self.config).get("use_weighted_sampler", False):
            self.trainer = WeightedTrainer(
                model=self.model,
                args=args,
                train_dataset=train_dataset,
                eval_dataset=val_dataset,
                compute_metrics=self.compute_metrics,
                config=self.config
            )
        else:
            self.trainer = VideoMAETrainer(
                model=self.model,
                args=args,
                train_dataset=train_dataset,
                eval_dataset=val_dataset,
                compute_metrics=self.compute_metrics,
                config=self.config
            )

        self.trainer.train()
        #############
        train_metrics = self.hf_trainer.evaluate(train_dataset, metric_key_prefix="train")
        logging.info(f"TRAIN METRICS: {train_metrics}")
        #############

    def infer(self, test_dataset, rank=0, world_size=1):
        if not self.config.MODEL:
            args = TrainingArguments(
            output_dir=self.save_dir,  # any directory, not used here
            per_device_eval_batch_size=1,
            )

            self.hf_trainer = HFTrainer(
                model=self.model,
                args=args,
                compute_metrics=self.compute_metrics  # optional, can compute later manually
            )

            preds_output = self.hf_trainer.predict(test_dataset)
            logits = preds_output.predictions
            if isinstance(logits, tuple):
                logits = logits[0]

            probs = torch.softmax(torch.tensor(logits), dim=-1).cpu().numpy()
            pred_idx = np.argmax(probs, axis=-1)
            confs = probs.max(axis=-1)

            class_names = test_dataset.label_map
            sample_ids = [item.get("id", str(i)) for i, item in enumerate(test_dataset.samples)]

            submission = {
                "version": "2.0",
                "task": "action_classification",
                "date": datetime.now().strftime("%Y-%m-%d"),
                "metadata": {"type": "predictions"},
                "data": [],
            }

            for sid, label_idx, conf in zip(sample_ids, pred_idx, confs):
                submission["data"].append({
                    "id": sid,
                    "labels": {
                        "action": {
                            "label": class_names[int(label_idx)],
                            "confidence": float(conf),
                        }
                    },
                })

            out_dir = os.path.join(self.save_dir, "final")
            os.makedirs(out_dir, exist_ok=True)
            out_path = os.path.join(out_dir, "predictions_test_epoch_final.json")
            with open(out_path, "w", encoding="utf-8") as f:
                json.dump(submission, f, indent=2)
            self.predictions_payload = submission
            self.predictions_path = out_path
            return submission

        else:
            from opensportslib.core.loss.builder import build_criterion
            from opensportslib.core.optimizer.builder import build_optimizer
            from opensportslib.core.scheduler.builder import build_scheduler
            from opensportslib.core.utils.data import tracking_collate_fn
            from torch.nn.parallel import DistributedDataParallel as DDP
            from torch.utils.data.distributed import DistributedSampler

            is_ddp = world_size > 1

            if is_ddp:
                torch.cuda.set_device(rank)
                self.device = torch.device(f"cuda:{rank}")
            else:
                self.device = select_device(self.config.SYSTEM)

            # model
            self.model = self.model.to(self.device)
            if is_ddp:
                self.model = DDP(self.model, device_ids=[rank]) 
                test_sampler = DistributedSampler(test_dataset, rank=rank, num_replicas=world_size)
            else:
                test_sampler = None

            modality = get_data_modality(self.config)
            is_tracking_modality = _is_tracking_graph_modality(modality)
            is_frames_modality = _is_frames_npy_modality(self.config, modality)
            collate_fn = tracking_collate_fn if is_tracking_modality else None
            test_dataloader_cfg = get_split_dataloader_cfg(self.config, "test")

            test_loader = DataLoader(
                test_dataset, 

                batch_size=test_dataloader_cfg.batch_size, 
                shuffle=False, 
                sampler=test_sampler,
                num_workers=getattr(test_dataloader_cfg, "num_workers", 0), 
                pin_memory=True,
                collate_fn=collate_fn
            )

            optimizer = self.optimizer if self.optimizer is not None else build_optimizer(self.model.parameters(), cfg=self.config.TRAIN.optimizer)
            scheduler = self.scheduler if self.scheduler is not None else build_scheduler(optimizer, cfg=self.config.TRAIN.scheduler)
            criterion = build_criterion(self.config.TRAIN.criterion)

            # Select trainer class based on modality
            if is_tracking_modality:
                TrainerClass = TrackingTrainerClassification
            elif is_frames_modality:
                TrainerClass = FramesTrainerClassification
            else:
                TrainerClass = MVTrainerClassification

            self.test_trainer = TrainerClass(
                train_loader=None,
                val_loader=None,
                test_loader=test_loader,
                model=self.model,
                optimizer=optimizer,
                scheduler=scheduler,
                criterion=criterion,
                class_weights=None,
                class_names=test_dataset.label_map,
                save_dir=self.save_dir,
                model_name=get_component_name_by_kind(self.config, "encoder") or "model",
                max_epochs=self.config.TRAIN.epochs,
                device=self.device,
                top_k=2,
                monitor=get_train_selection(self.config).get("monitor", "balanced_accuracy"),
                mode=get_train_selection(self.config).get("mode", "max"),
                revert_on_lr_reduction=(is_tracking_modality or is_frames_modality),
                config=self.config,
            )
            self.test_trainer.test(
                detailed_results=get_train_execution(self.config).get("detailed_results", False)
            )
            self.predictions_payload = getattr(
                self.test_trainer, "predictions_payload", None
            )
            self.predictions_path = getattr(
                self.test_trainer, "predictions_path", None
            )
            return self.predictions_payload

    def evaluate(self, pred_path, gt_path, class_names, exclude_labels=[]):

        label_to_idx = {v: k for k, v in class_names.items()}

        if isinstance(pred_path, dict):
            pred_data = pred_path
        elif isinstance(pred_path, str):
            with open(pred_path, encoding="utf-8") as f:
                pred_data = json.load(f)
        else:
            raise TypeError(
                f"Unsupported predictions type: {type(pred_path).__name__}. Expected dict or str."
            )

        with open(gt_path, encoding="utf-8") as f:
            gt_data = json.load(f)

        gt_dict = {}
        for item in gt_data["data"]:
            sid = item["id"]
            gt_label = item["labels"]["action"]["label"]
            if gt_label not in exclude_labels:
                if gt_label not in label_to_idx:
                    raise ValueError(
                        "Ground-truth label is not present in the resolved inference class list: "
                        f"{gt_label!r}."
                    )
                gt_dict[sid] = label_to_idx[gt_label]

        preds = []
        labels = []

        for item in pred_data["data"]:
            sid = item["id"]
            if sid not in gt_dict:
                continue

            pred_label = item["labels"]["action"]["label"]
            if pred_label not in label_to_idx:
                raise ValueError(
                    "Predicted label is not present in the resolved inference class list: "
                    f"{pred_label!r}."
                )

            preds.append(label_to_idx[pred_label])
            labels.append(gt_dict[sid])

        metrics = self.compute_metrics(
            (preds, labels),
            mode="labels"
        )
        return metrics


    def demo(self, model, video_paths):
        pass

    def save(self, model, path, processor=None, tokenizer=None, optimizer=None, epoch=None):
        """
        Save model checkpoint
        """
        save_checkpoint(model, path, processor, tokenizer, optimizer, epoch)
        logging.info(f"Model saved at {path}")

    def load(self, path, optimizer=None, scheduler=None):
        """
        Load model checkpoint. Returns loaded model, optimizer, epoch
        """
        from opensportslib.models.builder import build_model
        if self.model is None:
            self.model, _ = build_model(self.config, self.device)
        self.model, optimizer, scheduler, scaler, epoch, checkpoint = load_checkpoint(
            self.model, path, optimizer, scheduler, device=self.device
        )
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.scaler = scaler
        self.epoch = epoch
        logging.info(f"Model loaded from {path}, epoch: {epoch}")
        return self.model, self.optimizer, self.scheduler, self.epoch

compute_metrics(pred, mode='logits')

thin wrapper around the metric module.

Parameters:

Name	Type	Description	Default
pred		a tuple (logits, labels).	required
mode		"logits" or "labels" (default: "logits").	'logits'

Returns:

Type	Description
	a dictionary of classification metrics.

Source code in opensportslib/core/trainer/classification_trainer.py

def compute_metrics(self, pred, mode="logits"):
    """thin wrapper around the metric module.

    Args:
        pred: a tuple (logits, labels).
        mode: "logits" or "labels" (default: "logits").

    Returns:
        a dictionary of classification metrics.
    """
    return compute_classification_metrics(
        pred, top_k=2, mode=mode
    )

load(path, optimizer=None, scheduler=None)

Load model checkpoint. Returns loaded model, optimizer, epoch

Source code in opensportslib/core/trainer/classification_trainer.py

def load(self, path, optimizer=None, scheduler=None):
    """
    Load model checkpoint. Returns loaded model, optimizer, epoch
    """
    from opensportslib.models.builder import build_model
    if self.model is None:
        self.model, _ = build_model(self.config, self.device)
    self.model, optimizer, scheduler, scaler, epoch, checkpoint = load_checkpoint(
        self.model, path, optimizer, scheduler, device=self.device
    )
    self.optimizer = optimizer
    self.scheduler = scheduler
    self.scaler = scaler
    self.epoch = epoch
    logging.info(f"Model loaded from {path}, epoch: {epoch}")
    return self.model, self.optimizer, self.scheduler, self.epoch

save(model, path, processor=None, tokenizer=None, optimizer=None, epoch=None)

Save model checkpoint

Source code in opensportslib/core/trainer/classification_trainer.py

def save(self, model, path, processor=None, tokenizer=None, optimizer=None, epoch=None):
    """
    Save model checkpoint
    """
    save_checkpoint(model, path, processor, tokenizer, optimizer, epoch)
    logging.info(f"Model saved at {path}")

train(model, train_dataset, val_dataset=None, rank=0, world_size=1)

build all training components and run the loop.

detects the model type (HuggingFace vs. custom) and the data modality (video vs. tracking) to select the right trainer class, sampler, and collate function.

Parameters:

Name	Type	Description	Default
model		the classification model.	required
train_dataset		training ClassificationDataset.	required
val_dataset		validation ClassificationDataset (optional).	None
rank		GPU rank (0-indexed).	0
world_size		total number of GPUs.	1

Source code in opensportslib/core/trainer/classification_trainer.py

def train(self, model, train_dataset, val_dataset=None, rank=0, world_size=1):
    """build all training components and run the loop.

    detects the model type (HuggingFace vs. custom) and the data 
    modality (video vs. tracking) to select the right trainer class,
    sampler, and collate function.

    Args:
        model: the classification model.
        train_dataset: training ClassificationDataset.
        val_dataset: validation ClassificationDataset (optional).
        rank: GPU rank (0-indexed).
        world_size: total number of GPUs.
    """
    from opensportslib.core.loss.builder import build_criterion
    from opensportslib.core.optimizer.builder import build_optimizer
    from opensportslib.core.scheduler.builder import build_scheduler
    from opensportslib.core.utils.data import tracking_collate_fn
    from torch.nn.parallel import DistributedDataParallel as DDP
    from torch.utils.data.distributed import DistributedSampler

    is_ddp = world_size > 1
    modality = get_data_modality(self.config)
    is_tracking_modality = _is_tracking_graph_modality(modality)
    is_frames_modality = _is_frames_npy_modality(self.config, modality)
    seed = get_system_seed(self.config)

    g = torch.Generator()
    g.manual_seed(seed)

    # HuggingFace-backed models use the HF Trainer path when routed there.

    if is_ddp:
            torch.cuda.set_device(rank)
            self.device = torch.device(f"cuda:{rank}")
    else:
            self.device = select_device(self.config.SYSTEM)

    self.model = model.to(self.device)

    if is_ddp:
        self.model = DDP(self.model, device_ids=[rank])

    # Build components
    optimizer = build_optimizer(
        self.model.parameters(), cfg=self.config.TRAIN.optimizer
    )
    scheduler = build_scheduler(
        optimizer, cfg=self.config.TRAIN.scheduler
    )
    criterion = build_criterion(self.config.TRAIN.criterion)
    train_sampling = get_train_sampling(self.config)
    train_selection = get_train_selection(self.config)
    train_checkpoint = get_train_checkpoint(self.config)

    # --- class weights for the loss ---
    if train_sampling.get("use_weighted_loss", False):
        class_weights = train_dataset.get_class_weights(
            num_classes=train_dataset.num_classes(), sqrt=True
        ).to(self.device)
    else:
        class_weights = None

    # tracking modality needs a customm collate that merges PyG
    # Data objects into a single batched graph per timestamp.
    collate_fn = tracking_collate_fn if is_tracking_modality else None

    # --- train sampler ---
    if train_sampling.get("use_weighted_sampler", False):
        sample_weights = train_dataset.get_sample_weights()

        samples_per_class = train_sampling.get("samples_per_class")
        if samples_per_class:
            num_classes = train_dataset.num_classes()
            num_samples = samples_per_class * num_classes
        else:
            num_samples = len(sample_weights)

        if is_ddp:
            train_sampler = DistributedWeightedSampler(
                weights=sample_weights,
                num_replicas=world_size,
                rank=rank,
                replacement=True,
                num_samples=num_samples,
                seed=seed
            )
        else:
            train_sampler = WeightedRandomSampler(
                weights=sample_weights,
                num_samples=num_samples,
                replacement=True,
                generator=g
            )

        shuffle = False

    else:
        if is_ddp:
            train_sampler = DistributedSampler(
                train_dataset,
                num_replicas=world_size,
                rank=rank,
                shuffle=True,
                drop_last=True
            )
        else:
            train_sampler = None

        shuffle = not is_ddp


    # --- validation sampler ---
    if is_ddp:
        val_sampler = DistributedSampler(
            val_dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=False,
            drop_last=False
        )
    else:
        val_sampler = None

    train_dataloader_cfg = get_split_dataloader_cfg(self.config, "train")
    train_num_workers = getattr(train_dataloader_cfg, "num_workers", 0)
    train_pin_memory = getattr(train_dataloader_cfg, "pin_memory", self.device.type == "cuda")
    train_mp_context = getattr(train_dataloader_cfg, "mp_context", None)
    train_persistent_workers = getattr(train_dataloader_cfg, "persistent_workers", train_num_workers > 0)
    train_prefetch_factor = getattr(train_dataloader_cfg, "prefetch_factor", 4 if train_num_workers > 0 else None)

    if train_mp_context is not None:
        train_mp_context = mp.get_context(train_mp_context)

    train_loader = DataLoader(
        train_dataset,
        batch_size=train_dataloader_cfg.batch_size,
        shuffle=(train_sampler is None and shuffle),
        sampler=train_sampler,
        num_workers=train_num_workers,
        pin_memory=train_pin_memory,
        collate_fn=collate_fn,
        worker_init_fn=seed_worker,
        generator=g,
        drop_last=True,
        multiprocessing_context=train_mp_context,
        persistent_workers=train_persistent_workers,
        prefetch_factor=train_prefetch_factor,
    )

    valid_dataloader_cfg = get_split_dataloader_cfg(self.config, "valid")
    valid_num_workers = getattr(valid_dataloader_cfg, "num_workers", 0)
    valid_pin_memory = getattr(valid_dataloader_cfg, "pin_memory", self.device.type == "cuda")
    valid_mp_context = getattr(valid_dataloader_cfg, "mp_context", None)
    valid_persistent_workers = getattr(valid_dataloader_cfg, "persistent_workers", valid_num_workers > 0)
    valid_prefetch_factor = getattr(valid_dataloader_cfg, "prefetch_factor", 4 if valid_num_workers > 0 else None)

    if valid_mp_context is not None:
        valid_mp_context = mp.get_context(valid_mp_context)

    val_loader = DataLoader(
        val_dataset,
        batch_size=valid_dataloader_cfg.batch_size,
        shuffle=False,
        sampler=val_sampler,   
        num_workers=valid_num_workers,
        pin_memory=valid_pin_memory,
        collate_fn=collate_fn,
        worker_init_fn=seed_worker,
        generator=g,
        multiprocessing_context=valid_mp_context,
        persistent_workers=valid_persistent_workers,
        prefetch_factor=valid_prefetch_factor,
    )

    # select the modality-specific trainer.
    if is_tracking_modality:
        TrainerClass = TrackingTrainerClassification
    elif is_frames_modality:
        TrainerClass = FramesTrainerClassification
    else:
        TrainerClass = MVTrainerClassification

    self.trainer = TrainerClass(
        train_loader=train_loader,
        val_loader=val_loader,
        test_loader=None,
        model=self.model,
        optimizer=optimizer,
        scheduler=scheduler,
        criterion=criterion,
        class_weights=class_weights,
        class_names=train_dataset.label_map,
        save_dir=self.save_dir,
        model_name=get_component_name_by_kind(self.config, "encoder") or "model",
        max_epochs=self.config.TRAIN.epochs,
        device=self.device,
        top_k=2,
        patience=train_selection.get("patience", 0),
        monitor=train_selection.get("monitor", "balanced_accuracy"),
        mode=train_selection.get("mode", "max"),
        revert_on_lr_reduction=(is_tracking_modality or is_frames_modality),
        config=self.config,
    )

    self.trainer.train(
        epoch_start=self.epoch,
        save_every=train_checkpoint.get("save_every", 1),
    )
    return getattr(self.trainer, "best_checkpoint_path", None)

Copyright 2022 James Hong, Haotian Zhang, Matthew Fisher, Michael Gharbi, Kayvon Fatahalian

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

Trainer

Bases: ABC

Source code in opensportslib/core/trainer/localization_trainer.py

class Trainer(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def train(self):
        pass

Trainer_pl

Bases: Trainer

Trainer class used for models that rely on lightning modules.

Parameters:

Name	Type	Description	Default
cfg	dict	Dict config. It should contain the key 'max_epochs' and the key 'GPU'.	required

Source code in opensportslib/core/trainer/localization_trainer.py

class Trainer_pl(Trainer):
    """Trainer class used for models that rely on lightning modules.

    Args:
        cfg (dict): Dict config. It should contain the key 'max_epochs' and the key 'GPU'.
    """

    def __init__(self, cfg, work_dir):
        from opensportslib.core.utils.lightning import CustomProgressBar, MyCallback
        from pytorch_lightning.callbacks import ModelCheckpoint
        import pytorch_lightning as pl
        from pytorch_lightning.loggers import WandbLogger

        wandb_logger = None
        if wandb.run is not None:  # means init_wandb already ran
            wandb_logger = WandbLogger(experiment=wandb.run)

        self.work_dir = work_dir
        call = MyCallback()
        checkpoint_callback = ModelCheckpoint(
            dirpath=self.work_dir,
            filename="best",
            monitor=getattr(cfg.TRAIN, 'monitor', 'valid_loss'),   # or your metric
            mode=getattr(cfg.TRAIN, 'mode', 'min'),
            save_top_k=1,
        )
        self.trainer = pl.Trainer(
            logger=wandb_logger,
            max_epochs=getattr(cfg.TRAIN, 'max_epochs', 200),
            devices="auto",
            accelerator="auto",
            callbacks=[call, CustomProgressBar(refresh_rate=1), checkpoint_callback],
            num_sanity_val_steps=0,
        )
        self.best_checkpoint_path = None
        self.config = cfg

    def train(self, **kwargs):
        self.trainer.fit(**kwargs)

        # best_model = kwargs["model"].best_state

        # logging.info("Done training")
        # logging.info("Best epoch: {}".format(best_model.get("epoch")))
        # best_path = os.path.join(self.work_dir, "model.pth.tar")
        # self.best_checkpoint_path = best_path
        # torch.save(best_model, best_path)

        # logging.info("Model saved")
        # logging.info(best_path)

        self.best_checkpoint_path = self.trainer.checkpoint_callback.best_model_path

        from pytorch_lightning.utilities.rank_zero import rank_zero_only

        @rank_zero_only
        def log():
            logging.info("Done training")
            logging.info(f"Best model saved at: {self.best_checkpoint_path}")

            # Save the config file uniformly inside the work_dir
            if hasattr(self, 'config') and self.config is not None:
                from opensportslib.core.utils.config import save_config
                config_path = os.path.join(self.work_dir, "config.yaml")
                save_config(self.config, config_path)
                logging.info(f"Saved config: {config_path}")

        log()

Trainer_e2e

Bases: Trainer

Trainer class used for the e2e model.

Parameters:

Name	Type	Description	Default
args	dict	Dict of config.	required
optimizer	Optimizer	The optimizer to update model parameters. Set to None if validation epoch.	required
scaler	GradScaler	The gradient scaler for mixed precision training.	required
lr_scheduler		The learning rate scheduler.	required
work_dir	string	The folder in which the different files will be saved.	required
dali	bool	Whether videos are processed with dali or opencv.	required
repartitions	List[int]	List of gpus used data processing. Default: None.	None
cfg_test	dict	Dict of config for the inference (testing purpose) and evaluation of the test split. Occurs once training is done. Default: None.	None
cfg_challenge	dict	Dict of config for the inference (testing purpose) of the challenge split. Occurs once training is done. Default: None.	required
cfg_valid_data_frames	dict	Dict of config for the inference (testing purpose) and evaluation of the valid split. Occurs through the epochs after a certain number of epochs only if the criterion for the valid split is 'map'. Default: None.	None

Source code in opensportslib/core/trainer/localization_trainer.py

class Trainer_e2e(Trainer):
    """Trainer class used for the e2e model.

    Args:
        args (dict): Dict of config.
        model.
        optimizer (torch.optim.Optimizer): The optimizer to update model parameters. Set to None if validation epoch.
        scaler (torch.cuda.amp.GradScaler): The gradient scaler for mixed precision training.
        lr_scheduler : The learning rate scheduler.
        work_dir (string): The folder in which the different files will be saved.
        dali (bool): Whether videos are processed with dali or opencv.
        repartitions (List[int]): List of gpus used data processing.
            Default: None.
        cfg_test (dict): Dict of config for the inference (testing purpose) and evaluation of the test split. Occurs once training is done.
            Default: None.
        cfg_challenge (dict): Dict of config for the inference (testing purpose) of the challenge split. Occurs once training is done.
            Default: None.
        cfg_valid_data_frames (dict): Dict of config for the inference (testing purpose) and evaluation of the valid split. Occurs through the epochs after a certain number of epochs only if the criterion for the valid split is 'map'.
            Default: None.
    """

    def __init__(
        self,
        args,
        model,
        optimizer,
        scaler,
        lr_scheduler,
        work_dir,
        dali,
        repartitions=None,
        cfg_test=None,
        #cfg_challenge=None,
        cfg_valid_data_frames=None,
        start_epoch=0
    ):
        self.config = args
        self.losses = []
        self.best_epoch = 0
        execution = get_train_execution(args)
        self.best_criterion_valid = 0 if execution.get("criterion_valid") == "map" else float("inf")

        self.num_epochs = get_train_epochs(args)
        self.epoch = start_epoch
        self.model = model

        self.optimizer = optimizer
        self.scaler = scaler
        self.lr_scheduler = lr_scheduler

        self.acc_grad_iter = execution.get("acc_grad_iter")

        self.start_valid_epoch = execution.get("start_valid_epoch")
        self.criterion_valid = execution.get("criterion_valid")
        self.valid_map_every = execution.get("valid_map_every")
        #self.save_dir = work_dir
        self.dali = dali

        self.repartitions = repartitions
        self.cfg_test = cfg_test
        #self.cfg_challenge = cfg_challenge
        self.cfg_valid_data_frames = cfg_valid_data_frames

        self.best_checkpoint_path = None

        self.save_dir = work_dir #os.path.join(work_dir, run_name, run_id)
        os.makedirs(self.save_dir, exist_ok=True)
        try:
            wandb.watch(self.model, log="gradients", log_freq=100)
        except Exception:
            pass

    def save_checkpoint(self, epoch, is_best=False):
        """Save checkpoint with training state."""
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scaler_state_dict': self.scaler.state_dict(),
            'lr_state_dict': self.lr_scheduler.state_dict(),
            'best_epoch': self.best_epoch,
            'best_criterion_valid': self.best_criterion_valid
        }

        os.makedirs(self.save_dir, exist_ok=True)
        # Save latest checkpoint
        # latest_path = os.path.join(self.save_dir, f"latest_checkpoint_{epoch:03d}.pt")
        # torch.save(checkpoint, latest_path)
        # logging.info(f"Latest checkpoint saved: {latest_path}")

        # # Remove previous latest checkpoint
        # for f in os.listdir(self.save_dir):
        #     if f.startswith("latest_checkpoint_") and f != os.path.basename(latest_path):
        #         os.remove(os.path.join(self.save_dir, f))

        latest_path = os.path.join(self.save_dir, "latest_checkpoint.pt")
        torch.save(checkpoint, latest_path)
        logging.info(f"Latest checkpoint saved: {latest_path}")

        # Save best checkpoint if needed
        # if is_best:
        #     # best_path = os.path.join(self.save_dir, f"best_checkpoint_{epoch:03d}.pt")
        #     best_path = os.path.join(self.save_dir, f"best_checkpoint_{epoch:03d}.pt")
        #     torch.save(checkpoint, best_path)
        #     logging.info(f"Best checkpoint saved: {best_path}")

        #     # Remove previous best checkpoint
        #     for f in os.listdir(self.save_dir):
        #         if f.startswith("best_checkpoint_") and f != os.path.basename(best_path):
        #             os.remove(os.path.join(self.save_dir, f))

        if is_best:
            best_path = os.path.join(self.save_dir, "best_checkpoint.pt")
            self.best_checkpoint_path = best_path
            torch.save(checkpoint, best_path)
            logging.info(f"Best checkpoint saved: {best_path}")

        if self.config is not None:
            from opensportslib.core.utils.config import save_config
            config_path = os.path.join(self.save_dir, "config.yaml")
            save_config(self.config, config_path)
            logging.info(f"Saved config: {config_path}")

    def train(self, train_loader, valid_loader, classes):
        """Training loop with checkpoint management."""
        if self.criterion_valid == "map":
            data_obj_valid = build_dataset(self.config, split="valid_data_frames")
            dataset_Valid_Frames =  data_obj_valid.building_dataset(
                data_obj_valid.cfg,
                None,
                {"repartitions": self.repartitions, "classes": classes},
            )

        for epoch in range(self.epoch, self.num_epochs):
            train_loss = self.model.epoch(
                train_loader,
                self.dali,
                self.optimizer,
                self.scaler,
                lr_scheduler=self.lr_scheduler,
                acc_grad_iter=self.acc_grad_iter,
            )

            valid_loss = self.model.epoch(
                valid_loader, self.dali, acc_grad_iter=self.acc_grad_iter
            )
            print(
                f"[Epoch {epoch+1}/{self.num_epochs}] Train loss: {train_loss:.5f} Valid loss: {valid_loss:.5f}"
            )
            logging.info(
                f"[Epoch {epoch+1}/{self.num_epochs}] Train loss: {train_loss:.5f} Valid loss: {valid_loss:.5f}"
            )

            valid_mAP = 0
            is_best = False

            if self.criterion_valid == "loss":
                if valid_loss < self.best_criterion_valid:
                    self.best_criterion_valid = valid_loss
                    self.best_epoch = epoch
                    is_best = True
                    print("New best epoch!")
            elif self.criterion_valid == "map":
                if epoch >= self.start_valid_epoch and epoch % self.valid_map_every == 0:
                    pred_file = None
                    if self.save_dir is not None:
                        pred_file = os.path.join(
                            self.save_dir, f"pred-valid_{epoch:03d}"
                        )
                        os.makedirs(self.save_dir, exist_ok=True)
                    valid_mAP = infer_and_process_predictions_e2e(
                        self.model,
                        self.dali,
                        dataset_Valid_Frames,
                        "VALID",
                        classes,
                        pred_file,
                        dataloader_params=self.cfg_valid_data_frames.dataloader,
                    )
                    if valid_mAP > self.best_criterion_valid:
                        self.best_criterion_valid = valid_mAP
                        self.best_epoch = epoch
                        is_best = True
                        print("New best epoch!")
            else:
                print("Unknown criterion:", self.criterion_valid)

            self.losses.append(
                {
                    "epoch": epoch,
                    "train": train_loss,
                    "valid": valid_loss,
                    "valid_mAP": valid_mAP,
                }
            )

            # ---------------- W&B LOG ----------------
            if wandb.run is not None:
                payload = {
                    "epoch": epoch + 1,
                    "train/loss": train_loss,
                    "valid/loss": valid_loss,
                    "lr": self.optimizer.param_groups[0]["lr"],
                    "best/mAP": self.best_criterion_valid if self.criterion_valid == "map" else None,
                    "best/loss": self.best_criterion_valid if self.criterion_valid == "loss" else None,
                }
                # Whole-match mAP only runs every valid_map_every epochs.
                # Logging a placeholder 0 on the other epochs drew a sawtooth
                # collapsing to zero between real measurements; omit the key
                # instead so the chart connects the points that exist.
                if valid_mAP:
                    payload["valid/mAP"] = valid_mAP
                wandb.log(payload)

            if self.save_dir is not None:
                os.makedirs(self.save_dir, exist_ok=True)
                store_json(
                    os.path.join(self.save_dir, "loss.json"),
                    self.losses,
                    pretty=True
                )
                self.save_checkpoint(epoch, is_best)

        logging.info(f"Training completed. Best epoch: {self.best_epoch}")

        if self.dali:
            train_loader.delete()
            valid_loader.delete()
            if self.criterion_valid == "map":
                dataset_Valid_Frames.delete()

        if self.save_dir is not None:
            self._run_final_evaluation(classes, eval_splits=["valid"])

    def _run_final_evaluation(self, classes, eval_splits=["valid", "test"]):
        from opensportslib.core.utils.checkpoint import load_checkpoint, localization_remap
        """Run final evaluation using best model."""
        # Load best model for evaluation
        best_checkpoint_path = os.path.join(
            self.save_dir, f"best_checkpoint.pt"
        )
        self.model._model, _, _, _, epoch, _ = load_checkpoint(model=self.model._model,
                                        path=best_checkpoint_path,
                                        key_remap_fn=localization_remap)
        logging.info(f"Loaded best model from epoch {self.best_epoch}")

        for split in eval_splits:
            if split == "valid":
                cfg_tmp = self.cfg_valid_data_frames
                split = "valid_data_frames"
            elif split == "test":
                cfg_tmp = self.cfg_test
            # elif split == "challenge":
            #     cfg_tmp = self.cfg_challenge

            if cfg_tmp is None:
                continue

            split_path = getattr(cfg_tmp, "path", None)
            if split_path is None:
                split_path = getattr(cfg_tmp, "annotation_path", None)
            if split_path is None or not os.path.exists(split_path):
                continue

            data_obj = build_dataset(self.config, split=split)
            split_data = data_obj.building_dataset(
                data_obj.cfg,
                None,
                {"repartitions": self.repartitions, "classes": classes},
            )
            split_data.print_info()

            pred_file = None
            if self.save_dir is not None:
                pred_file = os.path.join(
                    self.save_dir, f"pred-{split}_{self.best_epoch:03d}"
                )

            infer_and_process_predictions_e2e(
                self.model,
                self.dali,
                split_data,
                split.upper(),
                classes,
                pred_file,
                calc_stats=split != "challenge",
                dataloader_params=cfg_tmp.dataloader,
            )

            if self.dali:
                split_data.delete()

        logging.info(f"Final evaluation completed. Best epoch: {self.best_epoch}")

save_checkpoint(epoch, is_best=False)

Save checkpoint with training state.

Source code in opensportslib/core/trainer/localization_trainer.py

def save_checkpoint(self, epoch, is_best=False):
    """Save checkpoint with training state."""
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': self.model.state_dict(),
        'optimizer_state_dict': self.optimizer.state_dict(),
        'scaler_state_dict': self.scaler.state_dict(),
        'lr_state_dict': self.lr_scheduler.state_dict(),
        'best_epoch': self.best_epoch,
        'best_criterion_valid': self.best_criterion_valid
    }

    os.makedirs(self.save_dir, exist_ok=True)
    # Save latest checkpoint
    # latest_path = os.path.join(self.save_dir, f"latest_checkpoint_{epoch:03d}.pt")
    # torch.save(checkpoint, latest_path)
    # logging.info(f"Latest checkpoint saved: {latest_path}")

    # # Remove previous latest checkpoint
    # for f in os.listdir(self.save_dir):
    #     if f.startswith("latest_checkpoint_") and f != os.path.basename(latest_path):
    #         os.remove(os.path.join(self.save_dir, f))

    latest_path = os.path.join(self.save_dir, "latest_checkpoint.pt")
    torch.save(checkpoint, latest_path)
    logging.info(f"Latest checkpoint saved: {latest_path}")

    # Save best checkpoint if needed
    # if is_best:
    #     # best_path = os.path.join(self.save_dir, f"best_checkpoint_{epoch:03d}.pt")
    #     best_path = os.path.join(self.save_dir, f"best_checkpoint_{epoch:03d}.pt")
    #     torch.save(checkpoint, best_path)
    #     logging.info(f"Best checkpoint saved: {best_path}")

    #     # Remove previous best checkpoint
    #     for f in os.listdir(self.save_dir):
    #         if f.startswith("best_checkpoint_") and f != os.path.basename(best_path):
    #             os.remove(os.path.join(self.save_dir, f))

    if is_best:
        best_path = os.path.join(self.save_dir, "best_checkpoint.pt")
        self.best_checkpoint_path = best_path
        torch.save(checkpoint, best_path)
        logging.info(f"Best checkpoint saved: {best_path}")

    if self.config is not None:
        from opensportslib.core.utils.config import save_config
        config_path = os.path.join(self.save_dir, "config.yaml")
        save_config(self.config, config_path)
        logging.info(f"Saved config: {config_path}")

train(train_loader, valid_loader, classes)

Training loop with checkpoint management.

Source code in opensportslib/core/trainer/localization_trainer.py

def train(self, train_loader, valid_loader, classes):
    """Training loop with checkpoint management."""
    if self.criterion_valid == "map":
        data_obj_valid = build_dataset(self.config, split="valid_data_frames")
        dataset_Valid_Frames =  data_obj_valid.building_dataset(
            data_obj_valid.cfg,
            None,
            {"repartitions": self.repartitions, "classes": classes},
        )

    for epoch in range(self.epoch, self.num_epochs):
        train_loss = self.model.epoch(
            train_loader,
            self.dali,
            self.optimizer,
            self.scaler,
            lr_scheduler=self.lr_scheduler,
            acc_grad_iter=self.acc_grad_iter,
        )

        valid_loss = self.model.epoch(
            valid_loader, self.dali, acc_grad_iter=self.acc_grad_iter
        )
        print(
            f"[Epoch {epoch+1}/{self.num_epochs}] Train loss: {train_loss:.5f} Valid loss: {valid_loss:.5f}"
        )
        logging.info(
            f"[Epoch {epoch+1}/{self.num_epochs}] Train loss: {train_loss:.5f} Valid loss: {valid_loss:.5f}"
        )

        valid_mAP = 0
        is_best = False

        if self.criterion_valid == "loss":
            if valid_loss < self.best_criterion_valid:
                self.best_criterion_valid = valid_loss
                self.best_epoch = epoch
                is_best = True
                print("New best epoch!")
        elif self.criterion_valid == "map":
            if epoch >= self.start_valid_epoch and epoch % self.valid_map_every == 0:
                pred_file = None
                if self.save_dir is not None:
                    pred_file = os.path.join(
                        self.save_dir, f"pred-valid_{epoch:03d}"
                    )
                    os.makedirs(self.save_dir, exist_ok=True)
                valid_mAP = infer_and_process_predictions_e2e(
                    self.model,
                    self.dali,
                    dataset_Valid_Frames,
                    "VALID",
                    classes,
                    pred_file,
                    dataloader_params=self.cfg_valid_data_frames.dataloader,
                )
                if valid_mAP > self.best_criterion_valid:
                    self.best_criterion_valid = valid_mAP
                    self.best_epoch = epoch
                    is_best = True
                    print("New best epoch!")
        else:
            print("Unknown criterion:", self.criterion_valid)

        self.losses.append(
            {
                "epoch": epoch,
                "train": train_loss,
                "valid": valid_loss,
                "valid_mAP": valid_mAP,
            }
        )

        # ---------------- W&B LOG ----------------
        if wandb.run is not None:
            payload = {
                "epoch": epoch + 1,
                "train/loss": train_loss,
                "valid/loss": valid_loss,
                "lr": self.optimizer.param_groups[0]["lr"],
                "best/mAP": self.best_criterion_valid if self.criterion_valid == "map" else None,
                "best/loss": self.best_criterion_valid if self.criterion_valid == "loss" else None,
            }
            # Whole-match mAP only runs every valid_map_every epochs.
            # Logging a placeholder 0 on the other epochs drew a sawtooth
            # collapsing to zero between real measurements; omit the key
            # instead so the chart connects the points that exist.
            if valid_mAP:
                payload["valid/mAP"] = valid_mAP
            wandb.log(payload)

        if self.save_dir is not None:
            os.makedirs(self.save_dir, exist_ok=True)
            store_json(
                os.path.join(self.save_dir, "loss.json"),
                self.losses,
                pretty=True
            )
            self.save_checkpoint(epoch, is_best)

    logging.info(f"Training completed. Best epoch: {self.best_epoch}")

    if self.dali:
        train_loader.delete()
        valid_loader.delete()
        if self.criterion_valid == "map":
            dataset_Valid_Frames.delete()

    if self.save_dir is not None:
        self._run_final_evaluation(classes, eval_splits=["valid"])

Inferer

Source code in opensportslib/core/trainer/localization_trainer.py

class Inferer:
    def __init__(self, cfg, model, infer_Spotting):
        """Initialize the Inferer class.

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
            model: The model that will be used to infer.
            infer_Spotting: The method that is used to infer.
        """
        self.cfg_model = cfg
        self.model = model
        self.infer_Spotting=infer_Spotting
        self.dali = get_loader_backend(cfg) == "dali"

    def infer(self, cfg, data, dataloader=None):
        """Infer actions from data.

        Args:
            data : The data from which we will infer.
            dataloader : The dataloader for the test data.

        Returns:
            Dict containing predictions
        """
        if self.infer_Spotting=="infer_JSON":
            return self.infer_JSON(cfg, self.model, data, dataloader)
        elif self.infer_Spotting=="infer_SN":    
            return self.infer_SN(cfg, self.model, data, dataloader)
        elif self.infer_Spotting=="infer_E2E":
            return self.infer_E2E(cfg, self.model, data, dataloader)
        elif self.infer_Spotting=="infer_H5HeaderRule":
            return self.infer_H5HeaderRule(cfg, self.model, data, dataloader)
        raise ValueError(f"Unsupported infer_Spotting method: {self.infer_Spotting}")

    def infer_common(self, cfg, model, data, dataloader=None):
        """Infer actions from data using a given model.

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
            model: The model that will be used to infer.
            data : The data from which we will infer.

        Returns:
            Dict containing predictions
        """
        # Run Inference on Dataset
        from opensportslib.core.utils.lightning import CustomProgressBar, MyCallback
        import pytorch_lightning as pl
        from pytorch_lightning.loggers import WandbLogger

        wandb_logger = None
        if wandb.run is not None:  # means init_wandb already ran
            wandb_logger = WandbLogger(experiment=wandb.run)

        if get_system_path(cfg, "work_dir") is not None and dataloader is not None:

            evaluator = pl.Trainer(
                logger=wandb_logger,
                callbacks=[CustomProgressBar()],
                devices=1,
                accelerator="auto",
                num_sanity_val_steps=0,
            )
            evaluator.predict(model, dataloader)
            return model.json_data


    def infer_JSON(self, cfg, model, data, dataloader=None):
        """Infer actions from data using a given model for NetVlad/CALF methods

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
            model: The model that will be used to infer.
            data : The data from which we will infer.

        Returns:
            Dict containing predictions
        """
        return self.infer_common(cfg, model, data, dataloader)


    def infer_SN(self, cfg, model, data, dataloader=None):
        """Infer actions from data using a given model for the SNV2 data

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
            model: The model that will be used to infer.
            data : The data from which we will infer.

        Returns:
            Dict containing predictions
        """
        return self.infer_common(cfg, model, data, dataloader)


    def infer_E2E(self, cfg, model, data, dataloader=None):
        """Infer actions from data using a given model for the e2espot method.

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
            model: The model that will be used to infer.
            data : The data from which we will infer.

        Returns:
            Dict containing predictions
        """
        pred_file = None
        work_dir = get_system_path(cfg, "work_dir")
        cfg_testset = get_split_cfg(cfg, "test")
        if work_dir is not None:
            pred_file = os.path.join(work_dir, cfg_testset.results)

        predictions = infer_and_process_predictions_e2e(
            model,
            self.dali,
            data,
            "infer",
            get_data_classes(cfg),
            pred_file,
            True,
            get_split_dataloader_cfg(cfg, "test"),
            return_pred=True,
        )

        if pred_file is not None:
            pred_json_file = os.path.join(pred_file + ".json")
            pred_recall_file = os.path.join(pred_file + ".recall.json.gz")
            #logging.info("Predictions saved")
            #logging.info(pred_json_file)
            logging.info("High recall predictions saved")
            logging.info(pred_recall_file)

        return predictions

    def infer_H5HeaderRule(self, cfg, model, data, dataloader=None):
        """Run rule-based H5 header spotting without Lightning or weights."""
        del cfg, dataloader
        return model.predict(data)

__init__(cfg, model, infer_Spotting)

Initialize the Inferer class.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that will be used to infer.	required
infer_Spotting		The method that is used to infer.	required

Source code in opensportslib/core/trainer/localization_trainer.py

def __init__(self, cfg, model, infer_Spotting):
    """Initialize the Inferer class.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model: The model that will be used to infer.
        infer_Spotting: The method that is used to infer.
    """
    self.cfg_model = cfg
    self.model = model
    self.infer_Spotting=infer_Spotting
    self.dali = get_loader_backend(cfg) == "dali"

infer(cfg, data, dataloader=None)

Infer actions from data.

Parameters:

Name	Type	Description	Default
data		The data from which we will infer.	required
dataloader		The dataloader for the test data.	required

Returns:

Type	Description
	Dict containing predictions

Source code in opensportslib/core/trainer/localization_trainer.py

def infer(self, cfg, data, dataloader=None):
    """Infer actions from data.

    Args:
        data : The data from which we will infer.
        dataloader : The dataloader for the test data.

    Returns:
        Dict containing predictions
    """
    if self.infer_Spotting=="infer_JSON":
        return self.infer_JSON(cfg, self.model, data, dataloader)
    elif self.infer_Spotting=="infer_SN":    
        return self.infer_SN(cfg, self.model, data, dataloader)
    elif self.infer_Spotting=="infer_E2E":
        return self.infer_E2E(cfg, self.model, data, dataloader)
    elif self.infer_Spotting=="infer_H5HeaderRule":
        return self.infer_H5HeaderRule(cfg, self.model, data, dataloader)
    raise ValueError(f"Unsupported infer_Spotting method: {self.infer_Spotting}")

infer_E2E(cfg, model, data, dataloader=None)

Infer actions from data using a given model for the e2espot method.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that will be used to infer.	required
data		The data from which we will infer.	required

Returns:

Type	Description
	Dict containing predictions

Source code in opensportslib/core/trainer/localization_trainer.py

def infer_E2E(self, cfg, model, data, dataloader=None):
    """Infer actions from data using a given model for the e2espot method.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model: The model that will be used to infer.
        data : The data from which we will infer.

    Returns:
        Dict containing predictions
    """
    pred_file = None
    work_dir = get_system_path(cfg, "work_dir")
    cfg_testset = get_split_cfg(cfg, "test")
    if work_dir is not None:
        pred_file = os.path.join(work_dir, cfg_testset.results)

    predictions = infer_and_process_predictions_e2e(
        model,
        self.dali,
        data,
        "infer",
        get_data_classes(cfg),
        pred_file,
        True,
        get_split_dataloader_cfg(cfg, "test"),
        return_pred=True,
    )

    if pred_file is not None:
        pred_json_file = os.path.join(pred_file + ".json")
        pred_recall_file = os.path.join(pred_file + ".recall.json.gz")
        #logging.info("Predictions saved")
        #logging.info(pred_json_file)
        logging.info("High recall predictions saved")
        logging.info(pred_recall_file)

    return predictions

infer_H5HeaderRule(cfg, model, data, dataloader=None)

Run rule-based H5 header spotting without Lightning or weights.

Source code in opensportslib/core/trainer/localization_trainer.py

def infer_H5HeaderRule(self, cfg, model, data, dataloader=None):
    """Run rule-based H5 header spotting without Lightning or weights."""
    del cfg, dataloader
    return model.predict(data)

infer_JSON(cfg, model, data, dataloader=None)

Infer actions from data using a given model for NetVlad/CALF methods

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that will be used to infer.	required
data		The data from which we will infer.	required

Returns:

Type	Description
	Dict containing predictions

Source code in opensportslib/core/trainer/localization_trainer.py

def infer_JSON(self, cfg, model, data, dataloader=None):
    """Infer actions from data using a given model for NetVlad/CALF methods

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model: The model that will be used to infer.
        data : The data from which we will infer.

    Returns:
        Dict containing predictions
    """
    return self.infer_common(cfg, model, data, dataloader)

infer_SN(cfg, model, data, dataloader=None)

Infer actions from data using a given model for the SNV2 data

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that will be used to infer.	required
data		The data from which we will infer.	required

Returns:

Type	Description
	Dict containing predictions

Source code in opensportslib/core/trainer/localization_trainer.py

def infer_SN(self, cfg, model, data, dataloader=None):
    """Infer actions from data using a given model for the SNV2 data

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model: The model that will be used to infer.
        data : The data from which we will infer.

    Returns:
        Dict containing predictions
    """
    return self.infer_common(cfg, model, data, dataloader)

infer_common(cfg, model, data, dataloader=None)

Infer actions from data using a given model.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that will be used to infer.	required
data		The data from which we will infer.	required

Returns:

Type	Description
	Dict containing predictions

Source code in opensportslib/core/trainer/localization_trainer.py

def infer_common(self, cfg, model, data, dataloader=None):
    """Infer actions from data using a given model.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model: The model that will be used to infer.
        data : The data from which we will infer.

    Returns:
        Dict containing predictions
    """
    # Run Inference on Dataset
    from opensportslib.core.utils.lightning import CustomProgressBar, MyCallback
    import pytorch_lightning as pl
    from pytorch_lightning.loggers import WandbLogger

    wandb_logger = None
    if wandb.run is not None:  # means init_wandb already ran
        wandb_logger = WandbLogger(experiment=wandb.run)

    if get_system_path(cfg, "work_dir") is not None and dataloader is not None:

        evaluator = pl.Trainer(
            logger=wandb_logger,
            callbacks=[CustomProgressBar()],
            devices=1,
            accelerator="auto",
            num_sanity_val_steps=0,
        )
        evaluator.predict(model, dataloader)
        return model.json_data

Evaluator

Evaluator class that is used to make easier the process of evaluate since there is only one evaluate method that uses the evaluate_Spotting method.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict.	required
evaluate_Spotting	method	The method that is used to evaluate.	required

Source code in opensportslib/core/trainer/localization_trainer.py

class Evaluator:
    """Evaluator class that is used to make easier the process of evaluate since there is only
    one evaluate method that uses the evaluate_Spotting method.

    Args:
        cfg (dict): Config dict.
        evaluate_Spotting (method): The method that is used to evaluate.
    """

    def __init__(self, cfg, evaluate_Spotting):
        self.cfg = cfg
        self.extract_fps = getattr(cfg.DATA, "extract_fps", 2)
        self.evaluate_Spotting = evaluate_Spotting

    def evaluate(self, cfg_testset, json_gz_file=None):
        """Evaluate predictions.

        Args:
            cfg_testset (dict): Config dict that contains informations for the predictions.
        """
        if self.evaluate_Spotting == "evaluate_pred_JSON":
            return self.evaluate_pred_JSON(cfg_testset, get_system_path(self.cfg, "work_dir"), json_gz_file, metric=cfg_testset.metric)
        elif self.evaluate_Spotting == "evaluate_pred_SN":
            return self.evaluate_pred_SN(cfg_testset, get_system_path(self.cfg, "work_dir"), json_gz_file, metric=cfg_testset.metric)
        elif self.evaluate_Spotting == "evaluate_pred_E2E":
            return self.evaluate_pred_E2E(cfg_testset, get_system_path(self.cfg, "work_dir"), json_gz_file, metric=cfg_testset.metric)


    # def evaluate_common_JSON(self, cfg, results, metric):
    #     if cfg.path == None:
    #         return
    #     with open(cfg.path) as f:
    #         GT_data = json.load(f)

    #     print(results)
    #     pred_path_is_json = False
    #     if results.endswith(".json"):
    #         pred_path_is_json = True
    #         with open(results) as f:
    #             pred_data = json.load(f)

    #     targets_numpy = list()
    #     detections_numpy = list()
    #     closests_numpy = list()

    #     if "labels" in GT_data.keys():
    #         classes = GT_data["labels"]
    #     else:
    #         assert isinstance(cfg.classes, list) or os.path.isfile(cfg.classes)
    #         if isinstance(cfg.classes, list):
    #             classes = cfg.classes

    #     classes = sorted(classes)
    #     EVENT_DICTIONARY = {x: i for i, x in enumerate(classes)}
    #     INVERSE_EVENT_DICTIONARY = {i: x for i, x in enumerate(classes)}

    #     if "videos" in GT_data.keys():
    #         videos = GT_data["videos"]
    #     else:
    #         videos = [GT_data]

    #     for game in tqdm.tqdm(videos):
    #         print(game.keys())
    #         # fetch labels
    #         labels = game["annotations"]
    #         if not pred_path_is_json:
    #             try:
    #                 pred_file = os.path.join(results, os.path.splitext(game["path"])[0], "results_spotting.json")
    #                 print(pred_file)
    #                 with open(pred_file) as f:
    #                     pred_data = json.load(f)
    #             except FileNotFoundError:
    #                 continue
    #         predictions = pred_data["predictions"]
    #         # convert labels to dense vector
    #         dense_labels = label2vector(
    #             labels,
    #             num_classes=len(classes),
    #             EVENT_DICTIONARY=EVENT_DICTIONARY,
    #             framerate=(
    #                 pred_data["fps"] if "fps" in pred_data.keys() else self.extract_fps
    #             ),
    #         )
    #         print(dense_labels.shape)
    #         # convert predictions to vector
    #         dense_predictions = predictions2vector(
    #             predictions,
    #             vector_size=game["num_frames"] if "num_frames" in game.keys() else None,
    #             framerate=(
    #                 pred_data["fps"] if "fps" in pred_data.keys() else self.extract_fps
    #             ),
    #             num_classes=len(classes),
    #             EVENT_DICTIONARY=EVENT_DICTIONARY,
    #         )
    #         print(dense_predictions.shape)

    #         targets_numpy.append(dense_labels)
    #         detections_numpy.append(dense_predictions)

    #         closest_numpy = np.zeros(dense_labels.shape) - 1
    #         # Get the closest action index
    #         closests_numpy.append(get_closest_action_index(dense_labels, closest_numpy))

    #     if targets_numpy:
    #         return compute_performances_mAP(
    #             metric,
    #             targets_numpy,
    #             detections_numpy,
    #             closests_numpy,
    #             INVERSE_EVENT_DICTIONARY,
    #         )
    #     else:
    #         logging.warning("No predictions found for evaluation. Returning None.")
    #         return None



    def evaluate_common_JSON(self, cfg, results, metric):
        def _normalize_pred_key(path_value):
            if path_value is None:
                return "", ""
            raw = str(path_value).strip().replace("\\", "/")
            no_ext = os.path.splitext(raw)[0]
            return raw, no_ext

        gt_path = getattr(cfg, "annotation_path", None)
        if gt_path is None:
            gt_path = getattr(cfg, "path", None)
        if gt_path is None:
            return

        # --------------------------------------------------
        # LOAD GT
        # --------------------------------------------------
        with open(gt_path, encoding="utf-8") as f:
            GT_data = json.load(f)

        # --------------------------------------------------
        # LOAD PRED FILE (json / json.gz / folder)
        # --------------------------------------------------
        pred_data = None
        pred_path_is_file = results.endswith(".json") or results.endswith(".json.gz")

        if pred_path_is_file:
            pred_data = load_gz_json(results) if results.endswith(".gz") else load_json(results)

        # detect v2 prediction
        pred_is_v2 = isinstance(pred_data, dict) and pred_data is not None and "data" in pred_data
        # --------------------------------------------------
        # CLASSES
        # --------------------------------------------------
        if isinstance(GT_data.get("labels"), dict):
            classes = list(GT_data["labels"].values())[0]["labels"]
        elif "labels" in GT_data:
            classes = GT_data["labels"]
        else:
            classes = cfg.classes

        classes = sorted(classes)
        EVENT_DICTIONARY = {x: i for i, x in enumerate(classes)}
        INVERSE_EVENT_DICTIONARY = {i: x for i, x in enumerate(classes)}

        # --------------------------------------------------
        # GT VIDEOS
        # --------------------------------------------------
        if "videos" in GT_data:
            videos = GT_data["videos"]
            gt_is_v2 = False
        else:
            all_segments = [
                segment
                for record in GT_data["data"]
                for segment in expand_localization_intervals(record)
            ]
            videos = [
                segment
                for segment in all_segments
                if segment["annotation_status"] == "verified"
            ]
            gt_is_v2 = True
            logging.info(
                "Localization evaluation selected %d/%d verified logical "
                "segments.",
                len(videos),
                len(all_segments),
            )

        # --------------------------------------------------
        # BUILD PRED LOOKUP IF V2
        # --------------------------------------------------
        pred_lookup = {}
        if pred_is_v2:
            for item in pred_data["data"]:
                video_path = item["inputs"][0]["path"]
                raw_key, no_ext_key = _normalize_pred_key(video_path)
                if raw_key:
                    pred_lookup[raw_key] = item
                if no_ext_key:
                    pred_lookup[no_ext_key] = item

        targets_numpy = []
        detections_numpy = []
        closests_numpy = []
        # Rate the dense vectors below end up sampled at; the mAP tolerances
        # are in seconds and must be converted with this same rate.
        eval_framerate = self.extract_fps

        # ==================================================
        # LOOP
        # ==================================================
        for game in tqdm.tqdm(videos):

            # ---------------- GT ----------------
            if gt_is_v2:
                video_path = game["logical_path"]
                labels = [
                    {
                        "label": event.get("label"),
                        "gameTime": event.get("gameTime"),
                        "position": int(
                            event.get("position_ms", event.get("position"))
                        ),
                    }
                    for event in game.get("events", [])
                ]
            else:
                video_path = game["path"]
                labels = game["annotations"]

            # ---------------- PRED ----------------
            if pred_path_is_file:
                # ===== V2 PRED =====
                if pred_is_v2:
                    raw_key, no_ext_key = _normalize_pred_key(video_path)
                    item = pred_lookup.get(raw_key) or pred_lookup.get(no_ext_key)
                    if item is None:
                        fps = self.extract_fps
                        predictions = []
                    else:
                        fps = item["inputs"][0].get(
                            "fps", self.extract_fps
                        )
                        predictions = [
                            {
                               "label": e.get("label"),
                               "gameTime": e.get("gameTime"),
                               "confidence": e.get("confidence"),
                               "position": int(
                                   e.get("position_ms", e.get("position"))
                               ),
                               "frame": e.get("frame")
                            }
                            for e in item.get("events", [])
                        ]

                # ===== OLD PRED =====
                else:
                    if "predictions" not in pred_data:
                        continue

                    predictions = pred_data["predictions"]
                    fps = pred_data.get("fps", self.extract_fps)

            else:
                # ===== FOLDER MODE =====
                pred_file = os.path.join(results, os.path.splitext(video_path)[0], "results_spotting.json")

                if not os.path.exists(pred_file):
                    fps = self.extract_fps
                    predictions = []
                else:
                    with open(pred_file, encoding="utf-8") as f:
                        pred_data_local = json.load(f)

                    if "data" in pred_data_local:
                        # v2 file inside folder
                        item = pred_data_local["data"][0]
                        fps = item["inputs"][0].get(
                            "fps", self.extract_fps
                        )

                        predictions = [
                            {
                               "label": e.get("label"),
                               "gameTime": e.get("gameTime"),
                               "confidence": e.get("confidence"),
                               "position": int(
                                   e.get("position_ms", e.get("position"))
                               ),
                               "frame": e.get("frame")
                            }
                            for e in item.get("events", [])
                        ]
                    else:
                        predictions = pred_data_local["predictions"]
                        fps = pred_data_local.get("fps", self.extract_fps)

            # ---------------- VECTORS ----------------
            # Size the dense vectors from the actual content instead of the
            # 90-minute default: a match clock can run well past 90 min
            # (kick-off offset, stoppage, half-time gap) and anything beyond
            # the cap is clamped onto the final bin, silently merging events.
            positions_ms = [a["position"] for a in labels] + [
                p["position"] for p in predictions
            ]
            vector_size = int(fps * (max(positions_ms) / 1000)) + 2 if positions_ms else None
            eval_framerate = fps

            dense_labels = label2vector(
                labels,
                num_classes=len(classes),
                EVENT_DICTIONARY=EVENT_DICTIONARY,
                framerate=fps,
                vector_size=vector_size,
            )

            dense_predictions = predictions2vector(
                predictions,
                vector_size=vector_size,
                framerate=fps,
                num_classes=len(classes),
                EVENT_DICTIONARY=EVENT_DICTIONARY,
            )
            targets_numpy.append(dense_labels)
            detections_numpy.append(dense_predictions)

            closest_numpy = np.zeros(dense_labels.shape) - 1
            closests_numpy.append(get_closest_action_index(dense_labels, closest_numpy))

        # --------------------------------------------------
        # METRICS
        # --------------------------------------------------
        if targets_numpy:
            return compute_performances_mAP(
                metric,
                targets_numpy,
                detections_numpy,
                closests_numpy,
                INVERSE_EVENT_DICTIONARY,
                framerate=eval_framerate,
            )
        else:
            logging.warning("No predictions found.")
            return None

    def evaluate_pred_E2E(self, cfg, work_dir, pred_path, metric="loose"):
        """Evaluate predictions infered with E2E method and display performances.
        Args:
            cfg (dict): It should containt the keys; classes (list of classes), path (path of the groundtruth data).
            It should contain the key nms_window if evaluation of raw predictions. It should containt the key extract_fps if predictions file do not contain the fps at which the frames have been processed to infer.
            work_dir: The folder path under which the prediction files are stored.
            pred_path: The path for predictions files. It can be:
                - folder path (that contains predictions files)
                - file path (if raw prediction that needs to be processed first)
            metric (string): metric used to evaluate.
                In ["loose","tight","at1","at2","at3","at4","at5"].
                Default: "loose".

        Returns
            The different mAPs computed.
        """

        results = pred_path

        if os.path.isfile(results) and (
            results.endswith(".gz") or results.endswith(".json")
        ):
            pred = (load_gz_json if results.endswith(".gz") else load_json)(results)
            # --------------------------------------------------
            # SUPPORT NEW V2 FORMAT (dict)
            # --------------------------------------------------
            if isinstance(pred, dict) and "data" in pred:
                internal = []

                for item in pred["data"]:
                    video = item["inputs"][0]["path"]
                    fps = item["inputs"][0].get("fps", self.extract_fps)

                    events = []
                    for ev in item.get("events", []):
                        events.append({
                            "frame": ev.get("frame"),
                            "label": ev.get("label"),
                            "confidence": ev.get("confidence"),
                            "position": int(ev.get("position_ms")),
                            "gameTime": ev.get("gameTime"),
                        })

                    internal.append({
                        "video": video,
                        "fps": fps,
                        "events": events,
                    })

                pred = internal
            nms_window = cfg.nms_window
            if isinstance(pred, list):
                if nms_window > 0:
                    logging.info("Applying NMS: " + str(nms_window))
                    pred = non_maximum_supression(pred, nms_window)

                eval_dir = os.path.join(work_dir, pred_path.split(".gz")[0].split(".json")[0])
                only_one_file = store_eval_files_json(pred, eval_dir)
                logging.info("Done processing prediction files!")
                if only_one_file:
                    results = os.path.join(eval_dir, "results_spotting.json")
                else:
                    results = eval_dir
        return self.evaluate_common_JSON(cfg, results, metric)


    def evaluate_pred_JSON(self, cfg, work_dir, pred_path, metric="loose"):
        """Evaluate predictions infered with Json files and display performances.
        Args:
            cfg (dict): It should containt the key path (path of the groundtruth data). It should containt the key classes (list of classes) if the different classes are not in the ground truth data.
            work_dir: The folder path under which the prediction files are stored.
            pred_path: The path for predictions files. It can be:
                - folder path (that contains predictions files)
                - json file path if evaluate only one json file.
            metric (string): metric used to evaluate.
                In ["loose","tight","at1","at2","at3","at4","at5"].
                Default: "loose".

        Returns
            The different mAPs computed.
        """
        return self.evaluate_common_JSON(cfg, os.path.join(work_dir, pred_path), metric)


    def evaluate_pred_SN(self, cfg, work_dir, pred_path, metric="loose"):
        """Evaluate predictions infered using SNV2 splits and display performances. This method should be used only for SNV2 dataset.
        Args:
            cfg (dict): It should containt the key path (path of the groundtruth data). It should containt the key classes (list of classes) if the different classes are not in the ground truth data.
            work_dir: The folder path under which the prediction files are stored.
            pred_path: The path for predictions files.
            metric (string): metric used to evaluate.
                In ["loose","tight","at1","at2","at3","at4","at5"].
                Default: "loose".

        Returns
            The different mAPs computed.
        """
        from SoccerNet.Evaluation.utils import INVERSE_EVENT_DICTIONARY_V2
        from SoccerNet.Evaluation.ActionSpotting import evaluate
        # challenge sets to be tested on EvalAI
        if "challenge" in cfg.split:
            print("Visit eval.ai to evaluate performances on Challenge set")
            return None
        # GT_path = cfg.data_root
        pred_path = os.path.join(work_dir, pred_path)
        results = evaluate(
            SoccerNet_path=cfg.data_root,
            Predictions_path=pred_path,
            split=cfg.split,
            prediction_file="results_spotting.json",
            version=getattr(cfg, "version", 2),
            metric=metric,
        )
        rows = []
        for i in range(len(results["a_mAP_per_class"])):
            label = INVERSE_EVENT_DICTIONARY_V2[i]
            rows.append(
                (
                    label,
                    "{:0.2f}".format(results["a_mAP_per_class"][i] * 100),
                    "{:0.2f}".format(results["a_mAP_per_class_visible"][i] * 100),
                    "{:0.2f}".format(results["a_mAP_per_class_unshown"][i] * 100),
                )
            )
        rows.append(
            (
                "Average mAP",
                "{:0.2f}".format(results["a_mAP"] * 100),
                "{:0.2f}".format(results["a_mAP_visible"] * 100),
                "{:0.2f}".format(results["a_mAP_unshown"] * 100),
            )
        )
        header = ["", "Any", "Visible", "Unseen"]
        result = tabulate(rows, headers=header)
        from opensportslib.core.utils.wandb import log_table_wandb
        log_table_wandb(name=f"Final Scores (Metric : {metric})", rows=rows, headers=header)
        logging.info("Best Performance at end of training ")
        logging.info("Metric: " + metric)
        logging.info("\n" + result)
        # logging.info("a_mAP visibility all: " +  str(results["a_mAP"]))
        # logging.info("a_mAP visibility all per class: " +  str( results["a_mAP_per_class"]))

        return results

evaluate(cfg_testset, json_gz_file=None)

Evaluate predictions.

Parameters:

Name	Type	Description	Default
cfg_testset	dict	Config dict that contains informations for the predictions.	required

Source code in opensportslib/core/trainer/localization_trainer.py

def evaluate(self, cfg_testset, json_gz_file=None):
    """Evaluate predictions.

    Args:
        cfg_testset (dict): Config dict that contains informations for the predictions.
    """
    if self.evaluate_Spotting == "evaluate_pred_JSON":
        return self.evaluate_pred_JSON(cfg_testset, get_system_path(self.cfg, "work_dir"), json_gz_file, metric=cfg_testset.metric)
    elif self.evaluate_Spotting == "evaluate_pred_SN":
        return self.evaluate_pred_SN(cfg_testset, get_system_path(self.cfg, "work_dir"), json_gz_file, metric=cfg_testset.metric)
    elif self.evaluate_Spotting == "evaluate_pred_E2E":
        return self.evaluate_pred_E2E(cfg_testset, get_system_path(self.cfg, "work_dir"), json_gz_file, metric=cfg_testset.metric)

evaluate_pred_E2E(cfg, work_dir, pred_path, metric='loose')

Evaluate predictions infered with E2E method and display performances. Args: cfg (dict): It should containt the keys; classes (list of classes), path (path of the groundtruth data). It should contain the key nms_window if evaluation of raw predictions. It should containt the key extract_fps if predictions file do not contain the fps at which the frames have been processed to infer. work_dir: The folder path under which the prediction files are stored. pred_path: The path for predictions files. It can be: - folder path (that contains predictions files) - file path (if raw prediction that needs to be processed first) metric (string): metric used to evaluate. In ["loose","tight","at1","at2","at3","at4","at5"]. Default: "loose".

Returns The different mAPs computed.

Source code in opensportslib/core/trainer/localization_trainer.py

def evaluate_pred_E2E(self, cfg, work_dir, pred_path, metric="loose"):
    """Evaluate predictions infered with E2E method and display performances.
    Args:
        cfg (dict): It should containt the keys; classes (list of classes), path (path of the groundtruth data).
        It should contain the key nms_window if evaluation of raw predictions. It should containt the key extract_fps if predictions file do not contain the fps at which the frames have been processed to infer.
        work_dir: The folder path under which the prediction files are stored.
        pred_path: The path for predictions files. It can be:
            - folder path (that contains predictions files)
            - file path (if raw prediction that needs to be processed first)
        metric (string): metric used to evaluate.
            In ["loose","tight","at1","at2","at3","at4","at5"].
            Default: "loose".

    Returns
        The different mAPs computed.
    """

    results = pred_path

    if os.path.isfile(results) and (
        results.endswith(".gz") or results.endswith(".json")
    ):
        pred = (load_gz_json if results.endswith(".gz") else load_json)(results)
        # --------------------------------------------------
        # SUPPORT NEW V2 FORMAT (dict)
        # --------------------------------------------------
        if isinstance(pred, dict) and "data" in pred:
            internal = []

            for item in pred["data"]:
                video = item["inputs"][0]["path"]
                fps = item["inputs"][0].get("fps", self.extract_fps)

                events = []
                for ev in item.get("events", []):
                    events.append({
                        "frame": ev.get("frame"),
                        "label": ev.get("label"),
                        "confidence": ev.get("confidence"),
                        "position": int(ev.get("position_ms")),
                        "gameTime": ev.get("gameTime"),
                    })

                internal.append({
                    "video": video,
                    "fps": fps,
                    "events": events,
                })

            pred = internal
        nms_window = cfg.nms_window
        if isinstance(pred, list):
            if nms_window > 0:
                logging.info("Applying NMS: " + str(nms_window))
                pred = non_maximum_supression(pred, nms_window)

            eval_dir = os.path.join(work_dir, pred_path.split(".gz")[0].split(".json")[0])
            only_one_file = store_eval_files_json(pred, eval_dir)
            logging.info("Done processing prediction files!")
            if only_one_file:
                results = os.path.join(eval_dir, "results_spotting.json")
            else:
                results = eval_dir
    return self.evaluate_common_JSON(cfg, results, metric)

evaluate_pred_JSON(cfg, work_dir, pred_path, metric='loose')

Evaluate predictions infered with Json files and display performances. Args: cfg (dict): It should containt the key path (path of the groundtruth data). It should containt the key classes (list of classes) if the different classes are not in the ground truth data. work_dir: The folder path under which the prediction files are stored. pred_path: The path for predictions files. It can be: - folder path (that contains predictions files) - json file path if evaluate only one json file. metric (string): metric used to evaluate. In ["loose","tight","at1","at2","at3","at4","at5"]. Default: "loose".

Returns The different mAPs computed.

Source code in opensportslib/core/trainer/localization_trainer.py

def evaluate_pred_JSON(self, cfg, work_dir, pred_path, metric="loose"):
    """Evaluate predictions infered with Json files and display performances.
    Args:
        cfg (dict): It should containt the key path (path of the groundtruth data). It should containt the key classes (list of classes) if the different classes are not in the ground truth data.
        work_dir: The folder path under which the prediction files are stored.
        pred_path: The path for predictions files. It can be:
            - folder path (that contains predictions files)
            - json file path if evaluate only one json file.
        metric (string): metric used to evaluate.
            In ["loose","tight","at1","at2","at3","at4","at5"].
            Default: "loose".

    Returns
        The different mAPs computed.
    """
    return self.evaluate_common_JSON(cfg, os.path.join(work_dir, pred_path), metric)

evaluate_pred_SN(cfg, work_dir, pred_path, metric='loose')

Evaluate predictions infered using SNV2 splits and display performances. This method should be used only for SNV2 dataset. Args: cfg (dict): It should containt the key path (path of the groundtruth data). It should containt the key classes (list of classes) if the different classes are not in the ground truth data. work_dir: The folder path under which the prediction files are stored. pred_path: The path for predictions files. metric (string): metric used to evaluate. In ["loose","tight","at1","at2","at3","at4","at5"]. Default: "loose".

Returns The different mAPs computed.

Source code in opensportslib/core/trainer/localization_trainer.py

def evaluate_pred_SN(self, cfg, work_dir, pred_path, metric="loose"):
    """Evaluate predictions infered using SNV2 splits and display performances. This method should be used only for SNV2 dataset.
    Args:
        cfg (dict): It should containt the key path (path of the groundtruth data). It should containt the key classes (list of classes) if the different classes are not in the ground truth data.
        work_dir: The folder path under which the prediction files are stored.
        pred_path: The path for predictions files.
        metric (string): metric used to evaluate.
            In ["loose","tight","at1","at2","at3","at4","at5"].
            Default: "loose".

    Returns
        The different mAPs computed.
    """
    from SoccerNet.Evaluation.utils import INVERSE_EVENT_DICTIONARY_V2
    from SoccerNet.Evaluation.ActionSpotting import evaluate
    # challenge sets to be tested on EvalAI
    if "challenge" in cfg.split:
        print("Visit eval.ai to evaluate performances on Challenge set")
        return None
    # GT_path = cfg.data_root
    pred_path = os.path.join(work_dir, pred_path)
    results = evaluate(
        SoccerNet_path=cfg.data_root,
        Predictions_path=pred_path,
        split=cfg.split,
        prediction_file="results_spotting.json",
        version=getattr(cfg, "version", 2),
        metric=metric,
    )
    rows = []
    for i in range(len(results["a_mAP_per_class"])):
        label = INVERSE_EVENT_DICTIONARY_V2[i]
        rows.append(
            (
                label,
                "{:0.2f}".format(results["a_mAP_per_class"][i] * 100),
                "{:0.2f}".format(results["a_mAP_per_class_visible"][i] * 100),
                "{:0.2f}".format(results["a_mAP_per_class_unshown"][i] * 100),
            )
        )
    rows.append(
        (
            "Average mAP",
            "{:0.2f}".format(results["a_mAP"] * 100),
            "{:0.2f}".format(results["a_mAP_visible"] * 100),
            "{:0.2f}".format(results["a_mAP_unshown"] * 100),
        )
    )
    header = ["", "Any", "Visible", "Unseen"]
    result = tabulate(rows, headers=header)
    from opensportslib.core.utils.wandb import log_table_wandb
    log_table_wandb(name=f"Final Scores (Metric : {metric})", rows=rows, headers=header)
    logging.info("Best Performance at end of training ")
    logging.info("Metric: " + metric)
    logging.info("\n" + result)
    # logging.info("a_mAP visibility all: " +  str(results["a_mAP"]))
    # logging.info("a_mAP visibility all per class: " +  str( results["a_mAP_per_class"]))

    return results

build_trainer(cfg, model=None, default_args=None, resume_from=None)

Build a trainer from config dict.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that is used to train. Needed only if E2E method because training do not rely on pytorch lightning. Default: None.	required
default_args	dict | None	Default initialization arguments. Default: None.	None

Returns:

Name	Type	Description
evaluator		The constructed trainer.

Source code in opensportslib/core/trainer/localization_trainer.py

def build_trainer(cfg, model=None, default_args=None, resume_from=None):
    """Build a trainer from config dict.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model : The model that is used to train. Needed only if E2E method because training do not rely on pytorch lightning.
            Default: None.
        default_args (dict | None, optional): Default initialization arguments.
            Default: None.

    Returns:
        evaluator: The constructed trainer.
    """
    if cfg.TRAIN.trainer.type == "trainer_e2e":
        checkpoint_dir = default_args["work_dir"]
        start_epoch = 0
        logging.info(f"Checkpoint directory: {checkpoint_dir}")

        # Handle checkpoint loading
        if resume_from is not None:
            optimizer = resume_from["optimizer"]
            scheduler = resume_from["scheduler"]
            scaler = resume_from["scaler"]
            start_epoch = resume_from["epoch"] + 1
            # Check if we've already reached target epochs
            if start_epoch >= get_train_epochs(cfg):
                logging.error(f"Model already trained for {start_epoch} epochs")
                logging.error(f"Target epochs in config: {get_train_epochs(cfg)}")
                logging.error("Please increase num_epochs in config to continue training")
                raise ValueError("Need to increase num_epochs to continue training")

            logging.info(f"Will continue training from epoch {start_epoch} to {get_train_epochs(cfg)}")
        else:
            logging.info("Building optimizer...")
            optimizer, scaler = build_optimizer(model._get_params(), cfg.TRAIN.optimizer)
            logging.info("Building scheduler...")
            scheduler = build_scheduler(optimizer, cfg.TRAIN.scheduler, default_args)

        trainer = Trainer_e2e(
            cfg,
            model,
            optimizer,
            scaler,
            scheduler,
            default_args["work_dir"],
            default_args["dali"],
            default_args["repartitions"],
            default_args["cfg_test"],
            #default_args["cfg_challenge"],
            default_args["cfg_valid_data_frames"],
            start_epoch=start_epoch
        )

        # Load training history if resuming
        if resume_from is not None:
            trainer.best_epoch = resume_from.get('best_epoch', 0)
            trainer.best_criterion_valid = resume_from.get('best_criterion_valid', 
                0 if cfg.TRAIN.execution.criterion_valid == "map" else float("inf"))
            logging.info(f"Restored best epoch: {trainer.best_epoch}")

    else:
        trainer = Trainer_pl(cfg, default_args["work_dir"])


    return trainer

build_inferer(cfg, model, default_args=None)

Build a inferer from config dict.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
model		The model that will be used to infer.	required
default_args	dict | None	Default initialization arguments. Default: None.	None

Returns:

Name	Type	Description
inferer		The constructed inferer.

Source code in opensportslib/core/trainer/localization_trainer.py

def build_inferer(cfg, model, default_args=None):
    """Build a inferer from config dict.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        model: The model that will be used to infer.
        default_args (dict | None, optional): Default initialization arguments.
            Default: None.

    Returns:
        inferer: The constructed inferer.
    """

    runner_type = get_runner_type(cfg)
    if runner_type == "runner_JSON":
        return Inferer(cfg=cfg, model=model, infer_Spotting="infer_JSON")
    if runner_type == "runner_pooling":
        return Inferer(cfg=cfg, model=model, infer_Spotting="infer_SN")
    if runner_type == "runner_CALF":
        return Inferer(cfg=cfg, model=model, infer_Spotting="infer_SN")
    if runner_type == "runner_e2e":
        return Inferer(cfg=cfg, model=model, infer_Spotting="infer_E2E")
    if runner_type == "runner_h5_header_rule":
        return Inferer(cfg=cfg, model=model, infer_Spotting="infer_H5HeaderRule")
    raise ValueError(f"Unsupported localization runner type: {runner_type}")

build_evaluator(cfg, default_args=None)

Build a evaluator from config dict.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
default_args	dict | None	Default initialization arguments. Default: None.	None

Returns:

Name	Type	Description
evaluator		The constructed evaluator.

Source code in opensportslib/core/trainer/localization_trainer.py

def build_evaluator(cfg, default_args=None):
    """Build a evaluator from config dict.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        default_args (dict | None, optional): Default initialization arguments.
            Default: None.

    Returns:
        evaluator: The constructed evaluator.
    """
    runner_type = get_runner_type(cfg)
    if runner_type == "runner_JSON":
        evaluator = Evaluator(cfg=cfg, evaluate_Spotting="evaluate_pred_JSON")
    elif runner_type == "runner_pooling":
        evaluator = Evaluator(cfg=cfg, evaluate_Spotting="evaluate_pred_SN")
    elif runner_type == "runner_CALF":
        evaluator = Evaluator(cfg=cfg, evaluate_Spotting="evaluate_pred_SN")
    elif runner_type == "runner_e2e":
        evaluator = Evaluator(cfg=cfg, evaluate_Spotting="evaluate_pred_E2E")

    return evaluator

Trainer/infer/eval helpers for VQA task.

XVarsVideoChatGPTDataCollator

Batch tokenized X-VARS samples including visual feature tensors.

Source code in opensportslib/core/trainer/vqa_trainer.py

class XVarsVideoChatGPTDataCollator:
    """Batch tokenized X-VARS samples including visual feature tensors."""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def __call__(self, instances):
        out = {
            "input_ids": torch.tensor([x["input_ids"] for x in instances], dtype=torch.long),
            "attention_mask": torch.tensor([x["attention_mask"] for x in instances], dtype=torch.long),
            "labels": torch.tensor([x["labels"] for x in instances], dtype=torch.long),
            "video_spatio_temporal_features": torch.stack(
                [torch.as_tensor(x["video_spatio_temporal_features"], dtype=torch.float32) for x in instances],
                dim=0,
            ),
        }
        return out

XVarsVideoChatGPTTrainer

Minimal Trainer wrapper that passes the tokenizer into the X-VARS model.

Source code in opensportslib/core/trainer/vqa_trainer.py

class XVarsVideoChatGPTTrainer:
    """Minimal Trainer wrapper that passes the tokenizer into the X-VARS model."""

    def __init__(self, *, model, tokenizer, args, train_dataset, eval_dataset=None, data_collator=None, callbacks=None):
        from transformers import Trainer

        class _Trainer(Trainer):
            def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
                del kwargs
                outputs = model(**inputs, tokenizer=tokenizer)
                loss = outputs.loss
                return (loss, outputs) if return_outputs else loss

        self._trainer = _Trainer(
            model=model,
            args=args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            data_collator=data_collator or XVarsVideoChatGPTDataCollator(tokenizer),
            callbacks=callbacks,
        )
        self.model = model
        self.tokenizer = tokenizer

    def train(self):
        return self._trainer.train()

    def save_state(self):
        return self._trainer.save_state()

VQAXVarsVideoChatGPTSFTDataset

Tokenized multimodal SFT dataset for the true X-VARS backend.

Source code in opensportslib/core/trainer/vqa_trainer.py

class VQAXVarsVideoChatGPTSFTDataset:
    """Tokenized multimodal SFT dataset for the true X-VARS backend."""

    def __init__(self, dataset, *, config=None, tokenizer, prompt_cfg: dict[str, Any], sft_cfg: dict[str, Any], xvars_cfg: dict[str, Any]):
        self.rows = []
        if config is not None and is_xvars_videochatgpt_backend(config):
            token_len = get_xvars_train_video_token_len(config)
        else:
            token_len = int(prompt_cfg.get("video_token_len", xvars_cfg.get("video_token_len", sft_cfg.get("video_token_len", 356))))
        max_seq_length = int(sft_cfg.get("max_seq_length", 768))
        for sample in dataset:
            refs = sample.get("references") or []
            if not refs:
                continue
            features = sample.get("video_spatio_temporal_features")
            if features is None:
                continue
            features = validate_xvars_feature_tensor(
                torch.as_tensor(features, dtype=torch.float32),
                expected_tokens=token_len,
                context=f"X-VARS training features for sample '{sample.get('id', 'unknown')}'",
            )
            references = refs if str(sft_cfg.get("reference_mode", "all")).lower() == "all" else refs[:1]
            for reference in references:
                if not str(reference).strip():
                    continue
                row = build_vqa_sft_text(
                    sample,
                    config=config,
                    prompt_cfg=prompt_cfg,
                    sft_cfg={**sft_cfg, "include_video_tokens": True, "video_token_len": token_len},
                    reference=str(reference),
                )
                tokenized = self._tokenize_row(row, tokenizer=tokenizer, max_seq_length=max_seq_length)
                if tokenized is None:
                    continue
                tokenized["video_spatio_temporal_features"] = features
                tokenized["id"] = str(sample.get("id", ""))
                tokenized["question"] = str(sample.get("question", ""))
                tokenized["prompt"] = row["prompt"]
                self.rows.append(tokenized)

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, idx):
        return self.rows[idx]

    @staticmethod
    def _tokenize_row(row, *, tokenizer, max_seq_length: int):
        ignore_index = -100
        prompt = str(row.get("prompt", "")).strip()
        completion = str(row.get("completion", "")).strip()
        if not prompt or not completion:
            return None
        full_text = f"{prompt} {completion}".strip()
        enc_full = tokenizer(full_text, truncation=True, max_length=max_seq_length, padding="max_length")
        enc_prompt = tokenizer(prompt, truncation=True, max_length=max_seq_length, padding="max_length")
        input_ids = list(enc_full["input_ids"])
        attn = list(enc_full["attention_mask"])
        prompt_ids = list(enc_prompt["input_ids"])
        full_len = int(sum(attn))
        prompt_len = int(sum(enc_prompt["attention_mask"]))
        if full_len <= prompt_len or input_ids[:prompt_len] != prompt_ids[:prompt_len]:
            return None
        eos_token = str(getattr(tokenizer, "eos_token", "") or "")
        eos_token_id = getattr(tokenizer, "eos_token_id", None)
        if completion.endswith(eos_token) and eos_token and eos_token_id is not None:
            input_ids[full_len - 1] = int(eos_token_id)
        labels = list(input_ids)
        for i in range(len(labels)):
            if i < prompt_len or attn[i] == 0:
                labels[i] = ignore_index
        if all(x == ignore_index for x in labels):
            return None
        return {"input_ids": input_ids, "attention_mask": attn, "labels": labels}

VQANativeQwenVLSFTDataset

Lazy tokenized multimodal SFT dataset for native Qwen VL backends.

Source code in opensportslib/core/trainer/vqa_trainer.py

class VQANativeQwenVLSFTDataset:
    """Lazy tokenized multimodal SFT dataset for native Qwen VL backends."""

    def __init__(
        self,
        dataset,
        *,
        model: QwenVLNativeModel,
        prompt_cfg: dict[str, Any],
        sft_cfg: dict[str, Any],
        cache_dir: str | None = None,
        split_name: str = "train",
        invalid_row_report_path: str | None = None,
        fail_on_invalid: bool = False,
        rank: int = 0,
    ):
        self.rows = []
        self.entries = []
        self._memory_cache = OrderedDict()
        self._cache_hits = 0
        self._cache_misses = 0
        self._materialized = 0
        self._invalid_rows: list[dict[str, Any]] = []
        self._invalid_row_keys: set[tuple[str, int]] = set()
        self.model = model
        self.prompt_cfg = dict(prompt_cfg or {})
        self.sft_cfg = dict(sft_cfg or {})
        self.cache_dir = str(cache_dir).strip() if cache_dir else ""
        self.split_name = str(split_name or "train")
        self.invalid_row_report_path = str(invalid_row_report_path or "").strip()
        self.fail_on_invalid = bool(fail_on_invalid)
        self.rank = int(rank)
        self.memory_cache_rows = max(0, int(self.sft_cfg.get("memory_cache_rows", 0) or 0))
        max_seq_length = int(sft_cfg.get("max_seq_length", 1024))
        disable_tqdm = bool(sft_cfg.get("disable_tqdm", True))
        self.max_seq_length = max_seq_length
        if self.cache_dir:
            os.makedirs(self.cache_dir, exist_ok=True)
        try:
            total = len(dataset)
        except Exception:
            total = None
        logging.info(
            "Starting native Qwen VL SFT manifest build | split=%s | source_samples=%s | max_seq_length=%s | "
            "disable_tqdm=%s | cache_dir=%s",
            self.split_name,
            total if total is not None else "unknown",
            max_seq_length,
            disable_tqdm,
            self.cache_dir or "disabled",
        )
        visual_samples = 0
        question_answer_pairs = 0
        iterator = tqdm(
            dataset,
            total=total,
            desc="Qwen VL SFT manifest",
            unit="sample",
            disable=disable_tqdm,
            leave=False,
        )
        for sample in iterator:
            if not disable_tqdm:
                iterator.set_postfix_str(f"id={sample.get('id')}")
            refs = sample.get("references") or []
            if not refs:
                continue
            if not (sample.get("video_path") or sample.get("frame_paths") or sample.get("video_frames")):
                continue
            visual_samples += 1
            references = refs if str(sft_cfg.get("reference_mode", "all")).lower() == "all" else refs[:1]
            for reference_idx, reference in enumerate(references):
                answer_text = str(reference).strip()
                if not answer_text:
                    continue
                question_answer_pairs += 1
                self.entries.append(
                    {
                        "sample": sample,
                        "answer_text": answer_text,
                        "reference_index": int(reference_idx),
                        "cache_path": self._cache_path(sample, answer_text, reference_idx),
                    }
                )
        logging.info(
            "Finished native Qwen VL SFT manifest build | split=%s | source_samples=%s | usable_visual_samples=%s | "
            "question_answer_pairs=%s | tokenized_rows=%s | cache_dir=%s",
            self.split_name,
            total if total is not None else "unknown",
            visual_samples,
            question_answer_pairs,
            len(self.entries),
            self.cache_dir or "disabled",
        )

    def __len__(self):
        return len(self.entries)

    def __getitem__(self, idx):
        total = len(self.entries)
        for offset in range(total):
            current_idx = (idx + offset) % total
            if current_idx in self._memory_cache:
                row = self._memory_cache.pop(current_idx)
                self._memory_cache[current_idx] = row
                return row

            entry = self.entries[current_idx]
            cache_path = entry.get("cache_path") or ""
            if cache_path and os.path.isfile(cache_path):
                row = torch.load(cache_path, map_location="cpu", weights_only=False)
                self._cache_hits += 1
                self._remember_row(current_idx, row)
                return row

            sample = entry["sample"]
            try:
                try:
                    row = self.model.build_training_inputs(
                        sample,
                        prompt_cfg=self.prompt_cfg,
                        answer_text=entry["answer_text"],
                        max_seq_length=self.max_seq_length,
                        reference_index=int(entry.get("reference_index", 0)),
                    )
                except TypeError as exc:
                    if "reference_index" not in str(exc):
                        raise
                    row = self.model.build_training_inputs(
                        sample,
                        prompt_cfg=self.prompt_cfg,
                        answer_text=entry["answer_text"],
                        max_seq_length=self.max_seq_length,
                    )
            except NativeQwenVLInvalidRowError as exc:
                invalid_record = self._record_invalid_row(sample, entry, exc)
                if self.fail_on_invalid:
                    raise RuntimeError(
                        "Invalid native Qwen VL training row encountered before backward/allreduce | "
                        f"rank={self.rank} | split={self.split_name} | id={invalid_record.get('sample_id')} | "
                        f"reference_index={invalid_record.get('reference_index')} | reason={invalid_record.get('reason')}"
                    ) from exc
                continue

            row["id"] = str(sample.get("id", ""))
            row["question"] = str(sample.get("question", ""))
            row["reference_index"] = int(entry.get("reference_index", 0))

            if cache_path:
                tmp_path = f"{cache_path}.tmp"
                torch.save(row, tmp_path)
                os.replace(tmp_path, cache_path)
            self._cache_misses += 1
            self._materialized += 1
            if self._materialized <= 3 or self._materialized % 100 == 0:
                logging.info(
                    "Materialized native Qwen VL SFT row | split=%s | built=%s/%s | cache_hits=%s | cache_misses=%s | id=%s",
                    self.split_name,
                    self._materialized,
                    len(self.entries),
                    self._cache_hits,
                    self._cache_misses,
                    row.get("id", ""),
                )
            self._remember_row(current_idx, row)
            return row
        raise ValueError(f"No valid native Qwen VL rows remained for split '{self.split_name}'.")

    def _remember_row(self, idx: int, row: dict[str, Any]) -> None:
        if self.memory_cache_rows <= 0:
            return
        self._memory_cache[idx] = row
        self._memory_cache.move_to_end(idx)
        while len(self._memory_cache) > self.memory_cache_rows:
            self._memory_cache.popitem(last=False)

    def _record_invalid_row(
        self,
        sample: dict[str, Any],
        entry: dict[str, Any],
        exc: NativeQwenVLInvalidRowError,
    ) -> dict[str, Any]:
        sample_id = str(sample.get("id", ""))
        reference_index = int(entry.get("reference_index", 0))
        key = (sample_id, reference_index)
        context = dict(getattr(exc, "context", {}) or {})
        record = {
            "split": self.split_name,
            "rank": self.rank,
            "sample_id": sample_id,
            "question": str(sample.get("question", "")),
            "reference_index": reference_index,
            "reason": str(exc),
            **context,
        }
        if key not in self._invalid_row_keys:
            self._invalid_row_keys.add(key)
            self._invalid_rows.append(record)
            logging.warning(
                "Skipping invalid native Qwen VL row | split=%s | rank=%s | id=%s | reference_index=%s | reason=%s",
                self.split_name,
                self.rank,
                sample_id,
                reference_index,
                record["reason"],
            )
            self._write_invalid_row_report()
        return record

    def _write_invalid_row_report(self) -> None:
        if not self.invalid_row_report_path:
            return
        os.makedirs(os.path.dirname(self.invalid_row_report_path), exist_ok=True)
        with open(self.invalid_row_report_path, "w", encoding="utf-8") as f:
            json.dump(
                {
                    "kind": "native_qwen_vl_invalid_rows",
                    "split": self.split_name,
                    "count": len(self._invalid_rows),
                    "rows": self._invalid_rows,
                },
                f,
                indent=2,
                ensure_ascii=True,
            )

    @property
    def invalid_row_count(self) -> int:
        return len(self._invalid_rows)

    def _cache_path(self, sample: dict[str, Any], answer_text: str, reference_idx: int) -> str:
        if not self.cache_dir:
            return ""
        payload = {
            "split": self.split_name,
            "sample_id": str(sample.get("id", "")),
            "question": str(sample.get("question", "")),
            "answer_text": str(answer_text),
            "reference_index": int(reference_idx),
            "video_path": self._file_fingerprint(sample.get("video_path")),
            "frame_paths": [self._file_fingerprint(path) for path in list(sample.get("frame_paths") or [])],
            "frame_count": len(list(sample.get("video_frames") or [])),
            "prompt_cfg": self.prompt_cfg,
            "max_seq_length": self.max_seq_length,
            "model_id": str(getattr(self.model, "model_id", "")),
            "visual_input_mode": str(getattr(self.model, "visual_input_mode", "")),
            "num_frames": int(getattr(self.model, "num_frames", 0) or 0),
        }
        digest = hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode("utf-8")).hexdigest()
        return os.path.join(self.cache_dir, f"{digest}.pt")

    @staticmethod
    def _file_fingerprint(path: Any) -> dict[str, Any]:
        raw = str(path or "").strip()
        if not raw:
            return {"path": "", "exists": False}
        info = {"path": raw, "exists": os.path.exists(raw)}
        if info["exists"]:
            try:
                stat = os.stat(raw)
                info["size"] = int(stat.st_size)
                info["mtime_ns"] = int(stat.st_mtime_ns)
            except OSError:
                pass
        return info

VQAXVarsVideoChatGPTLoraTrainer

True X-VARS LoRA trainer preserving multimodal feature tensors.

Source code in opensportslib/core/trainer/vqa_trainer.py

class VQAXVarsVideoChatGPTLoraTrainer:
    """True X-VARS LoRA trainer preserving multimodal feature tensors."""

    def __init__(self, config):
        self.config = config

    def train(
        self,
        train_data,
        valid_data=None,
        *,
        rank: int = 0,
        world_size: int = 1,
        use_wandb: bool = False,
    ) -> str:
        execution = get_train_execution(self.config)
        prompt_cfg = _as_dict(execution.get("prompt"))
        sft_cfg = _as_dict(execution.get("sft"))
        hf_cfg = _as_dict(execution.get("hf"))
        lora_cfg = _as_dict(execution.get("lora"))
        quant_cfg = _as_dict(execution.get("quantization"))
        checkpoint_cfg = _as_dict(execution.get("checkpoint"))
        xvars_cfg = _as_dict(execution.get("xvars"))
        generated_validation_cfg = _as_dict(execution.get("generated_validation"))

        save_root = get_system_path(self.config, "save_dir", "./checkpoints") or "./checkpoints"
        output_dir = os.path.join(save_root, "xvars_videochatgpt_lora")
        os.makedirs(output_dir, exist_ok=True)

        prompt_token_len = _resolve_vqa_video_token_len(self.config, prompt_cfg=prompt_cfg, sft_cfg=sft_cfg)
        model_id = get_xvars_train_model_id(self.config, default=str(hf_cfg.get("model_id", "base_model_videoChatGPT")))
        metadata = {
            "backend": "xvars_videochatgpt_lora",
            "model_id": model_id,
            "status": "metadata_only",
            "multimodal_training": True,
            "video_token_len": prompt_token_len,
            "lora_target_modules": list(lora_cfg.get("target_modules") or DEFAULT_XVARS_TARGET_MODULES),
            "num_train_samples": 0,
            "num_valid_samples": 0,
        }
        if bool(execution.get("dry_run", False)):
            # Dry-run still validates that dataset items carry visual features.
            metadata["num_train_samples"] = sum(1 for x in train_data if x.get("video_spatio_temporal_features") is not None)
            if valid_data is not None:
                metadata["num_valid_samples"] = sum(1 for x in valid_data if x.get("video_spatio_temporal_features") is not None)
            return self._write_artifacts(output_dir, metadata)

        require_optional_package("transformers", "pip install transformers")
        require_optional_package("peft", "pip install peft")
        from transformers import AutoTokenizer, TrainingArguments

        local_files_only = bool(hf_cfg.get("local_files_only", False))
        tokenizer = AutoTokenizer.from_pretrained(
            get_xvars_train_tokenizer_id(self.config),
            local_files_only=local_files_only,
            use_fast=False,
            model_max_length=int(hf_cfg.get("model_max_length", 1048)),
            padding_side=str(hf_cfg.get("padding_side", "right")),
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        bnb_config = build_bitsandbytes_config(quant_cfg)
        model_kwargs = {"local_files_only": local_files_only}
        cuda_device_index = get_hf_cuda_device_index(self.config, hf_cfg)
        if bnb_config is not None:
            model_kwargs["quantization_config"] = bnb_config
            if get_hf_prefer_cuda(self.config, hf_cfg) and torch.cuda.is_available():
                if cuda_device_index is not None:
                    torch.cuda.set_device(cuda_device_index)
                    model_kwargs["device_map"] = {"": cuda_device_index}
                else:
                    model_kwargs["device_map"] = {"": torch.cuda.current_device()}
        base_lm = load_videochatgpt_compatible_causal_lm(model_id, **model_kwargs)
        from opensportslib.core.utils.hf_runtime import _ensure_video_special_tokens

        _ensure_video_special_tokens(tokenizer, base_lm)
        model = XVarsVideoChatGPTCausalLM.from_pretrained_projector(
            base_lm,
            xvars_cfg.get("projection_path"),
            mm_hidden_size=get_vqa_mm_hidden_size(self.config, default=1024),
        )
        model.config.use_cache = False
        lora_cfg = dict(lora_cfg)
        lora_cfg.setdefault("target_modules", DEFAULT_XVARS_TARGET_MODULES)
        model = apply_lora_for_causal_lm(model, lora_cfg, distributed=int(world_size) > 1)
        model.config.use_cache = False

        logging.info("Building X-VARS SFT datasets | rank=%s", rank)
        train_sft = VQAXVarsVideoChatGPTSFTDataset(
            train_data,
            config=self.config,
            tokenizer=tokenizer,
            prompt_cfg=prompt_cfg,
            sft_cfg=sft_cfg,
            xvars_cfg=xvars_cfg,
        )
        valid_sft = (
            VQAXVarsVideoChatGPTSFTDataset(
                valid_data,
                config=self.config,
                tokenizer=tokenizer,
                prompt_cfg=prompt_cfg,
                sft_cfg=sft_cfg,
                xvars_cfg=xvars_cfg,
            )
            if valid_data is not None
            else None
        )
        if len(train_sft) == 0:
            raise ValueError("No multimodal X-VARS training rows were usable.")
        metadata["num_train_samples"] = len(train_sft)
        metadata["num_valid_samples"] = len(valid_sft) if valid_sft is not None else 0
        feature_shape = tuple(train_sft[0]["video_spatio_temporal_features"].shape) if len(train_sft) else None
        logging.info(
            "Built X-VARS SFT datasets | rank=%s | train=%s | valid=%s | feature_shape=%s",
            rank,
            metadata["num_train_samples"],
            metadata["num_valid_samples"],
            feature_shape,
        )
        train_bs, eval_bs = _resolve_sft_per_device_batch_sizes(self.config, sft_cfg)

        optimizer_cfg = get_train_optimizer(self.config)
        execution_cfg = get_train_execution(self.config)
        fp16, bf16 = _resolve_training_precision_flags(self.config, sft_cfg)
        eval_strategy_key = "evaluation_strategy"
        try:
            if "eval_strategy" in inspect.signature(TrainingArguments.__init__).parameters:
                eval_strategy_key = "eval_strategy"
        except Exception:
            eval_strategy_key = "evaluation_strategy"

        training_kwargs = {
            "output_dir": output_dir,
            "per_device_train_batch_size": train_bs,
            "per_device_eval_batch_size": eval_bs,
            "gradient_accumulation_steps": int(execution_cfg.get("acc_grad_iter", sft_cfg.get("gradient_accumulation_steps", 1))),
            "num_train_epochs": float(get_train_epochs(self.config) or sft_cfg.get("num_train_epochs", 1)),
            "learning_rate": float(optimizer_cfg.get("lr", sft_cfg.get("learning_rate", 2e-4))),
            "logging_steps": int(execution_cfg.get("log_interval", sft_cfg.get("logging_steps", 1))),
            "optim": str(optimizer_cfg.get("hf_optim", "paged_adamw_8bit")),
            "weight_decay": float(optimizer_cfg.get("weight_decay", 0.001)),
            "dataloader_num_workers": int(sft_cfg.get("dataloader_num_workers", 0)),
            "seed": get_system_seed(self.config),
            "data_seed": get_system_seed(self.config),
            "lr_scheduler_type": str(sft_cfg.get("lr_scheduler_type", "constant")),
            "save_strategy": str(sft_cfg.get("save_strategy", "epoch")),
            "report_to": ["wandb"] if use_wandb else [],
            "remove_unused_columns": False,
            "disable_tqdm": bool(sft_cfg.get("disable_tqdm", True)),
            "use_cpu": not get_hf_prefer_cuda(self.config, hf_cfg),
            "fp16": fp16,
            "bf16": bf16,
            "gradient_checkpointing": bool(sft_cfg.get("gradient_checkpointing", False)),
        }
        max_steps = sft_cfg.get("max_steps")
        if max_steps is not None and int(max_steps) > 0:
            training_kwargs["max_steps"] = int(max_steps)
        if int(world_size) > 1:
            training_kwargs["ddp_find_unused_parameters"] = bool(sft_cfg.get("ddp_find_unused_parameters", True))
            training_kwargs["gradient_checkpointing_kwargs"] = {"use_reentrant": False}
            training_kwargs["local_rank"] = int(rank)
        training_kwargs[eval_strategy_key] = _normalize_eval_strategy(
            sft_cfg.get("evaluation_strategy", "epoch"),
            default="epoch",
        )
        args = TrainingArguments(**training_kwargs)
        generated_callback = None
        if bool(generated_validation_cfg.get("enabled", False)):
            generated_callback = _build_xvars_generated_validation_callback(
                tokenizer=tokenizer,
                rows=train_sft.rows,
                output_dir=output_dir,
                validation_cfg=generated_validation_cfg,
                use_step_schedule=bool(training_kwargs.get("max_steps", 0)),
            )
        trainer = XVarsVideoChatGPTTrainer(
            model=model,
            tokenizer=tokenizer,
            args=args,
            train_dataset=train_sft,
            eval_dataset=valid_sft,
            data_collator=XVarsVideoChatGPTDataCollator(tokenizer),
            callbacks=[generated_callback] if generated_callback is not None else None,
        )
        logging.info(
            "Starting X-VARS trainer.train | rank=%s | world_size=%s | precision=%s | "
            "max_seq_length=%s | lora_targets=%s | gradient_checkpointing=%s | "
            "ddp_find_unused_parameters=%s | use_cache=%s",
            rank,
            world_size,
            "fp16" if fp16 else "bf16" if bf16 else "fp32",
            int(sft_cfg.get("max_seq_length", 768)),
            list(lora_cfg["target_modules"]),
            training_kwargs["gradient_checkpointing"],
            training_kwargs.get("ddp_find_unused_parameters"),
            model.config.use_cache,
        )
        trainer.train()
        logging.info("Finished X-VARS trainer.train | rank=%s", rank)
        if rank == 0 and bool(checkpoint_cfg.get("save_adapter", True)):
            trainer.model.save_pretrained(output_dir)
            tokenizer.save_pretrained(output_dir)
        metadata["status"] = "trained"
        selected_output_dir = output_dir
        if rank == 0 and generated_callback is not None:
            generated_path = os.path.join(output_dir, "generated_validation.json")
            with open(generated_path, "w", encoding="utf-8") as f:
                json.dump(generated_callback.history, f, indent=2)
            accepted = any(record.get("accepted") for record in generated_callback.history)
            metadata["generated_validation_accepted"] = accepted
            metadata["generated_validation_history"] = generated_path
            if generated_callback.history:
                selected_output_dir = os.path.join(output_dir, "generated_validation_best")
                metadata["best_generated_checkpoint"] = selected_output_dir
        if rank == 0:
            self._write_artifacts(output_dir, metadata)
            if selected_output_dir != output_dir:
                self._write_artifacts(selected_output_dir, metadata)
            if bool(generated_validation_cfg.get("require_relevance", False)) and not metadata.get(
                "generated_validation_accepted", False
            ):
                raise RuntimeError(
                    "X-VARS generated-answer relevance gate failed. Inspect generated_validation.json before "
                    "starting or accepting a full training run."
                )
            return selected_output_dir
        return output_dir

    def _write_artifacts(self, output_dir: str, metadata: dict[str, Any]) -> str:
        return _save_training_artifacts(
            self.config,
            output_dir,
            metadata,
            marker_text="X-VARS VideoChatGPT LoRA adapter artifacts are stored in the parent directory when training runs.\n",
        )

VQAQwenXVarsLoraTrainer

Qwen LoRA trainer preserving the X-VARS multimodal feature contract.

Source code in opensportslib/core/trainer/vqa_trainer.py

class VQAQwenXVarsLoraTrainer:
    """Qwen LoRA trainer preserving the X-VARS multimodal feature contract."""

    def __init__(self, config):
        self.config = config

    def train(
        self,
        train_data,
        valid_data=None,
        *,
        rank: int = 0,
        world_size: int = 1,
        use_wandb: bool = False,
    ) -> str:
        execution = get_train_execution(self.config)
        prompt_cfg = _as_dict(execution.get("prompt"))
        sft_cfg = _as_dict(execution.get("sft"))
        hf_cfg = _as_dict(execution.get("hf"))
        lora_cfg = _as_dict(execution.get("lora"))
        quant_cfg = _as_dict(execution.get("quantization"))
        checkpoint_cfg = _as_dict(execution.get("checkpoint"))
        generated_validation_cfg = _as_dict(execution.get("generated_validation"))

        save_root = get_system_path(self.config, "save_dir", "./checkpoints") or "./checkpoints"
        output_dir = os.path.join(save_root, "qwen_xvars_lora")
        os.makedirs(output_dir, exist_ok=True)

        prompt_token_len = _resolve_vqa_video_token_len(self.config, prompt_cfg=prompt_cfg, sft_cfg=sft_cfg)
        model_id = get_xvars_train_model_id(self.config, default=str(hf_cfg.get("model_id", "Qwen/Qwen2.5-7B-Instruct")))
        metadata = {
            "backend": "qwen_xvars_lora",
            "model_id": model_id,
            "status": "metadata_only",
            "multimodal_training": True,
            "video_token_len": prompt_token_len,
            "lora_target_modules": list(lora_cfg.get("target_modules") or ["q_proj", "v_proj"]),
            "num_train_samples": 0,
            "num_valid_samples": 0,
        }
        if bool(execution.get("dry_run", False)):
            metadata["num_train_samples"] = sum(1 for x in train_data if x.get("video_spatio_temporal_features") is not None)
            if valid_data is not None:
                metadata["num_valid_samples"] = sum(1 for x in valid_data if x.get("video_spatio_temporal_features") is not None)
            return self._write_artifacts(output_dir, metadata)

        require_optional_package("transformers", "pip install transformers")
        require_optional_package("peft", "pip install peft")
        from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments

        local_files_only = bool(hf_cfg.get("local_files_only", False))
        tokenizer = AutoTokenizer.from_pretrained(
            get_xvars_train_tokenizer_id(self.config),
            local_files_only=local_files_only,
            use_fast=False,
            model_max_length=int(hf_cfg.get("model_max_length", 1048)),
            padding_side=str(hf_cfg.get("padding_side", "right")),
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        bnb_config = build_bitsandbytes_config(quant_cfg)
        model_kwargs = {"local_files_only": local_files_only}
        cuda_device_index = get_hf_cuda_device_index(self.config, hf_cfg)
        if bnb_config is not None:
            model_kwargs["quantization_config"] = bnb_config
            if get_hf_prefer_cuda(self.config, hf_cfg) and torch.cuda.is_available():
                if cuda_device_index is not None:
                    torch.cuda.set_device(cuda_device_index)
                    model_kwargs["device_map"] = {"": cuda_device_index}
                else:
                    model_kwargs["device_map"] = {"": torch.cuda.current_device()}
        base_lm = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
        from opensportslib.core.utils.hf_runtime import _ensure_video_special_tokens

        _ensure_video_special_tokens(tokenizer, base_lm)
        model = QwenXVarsCausalLM(base_lm, mm_hidden_size=get_vqa_mm_hidden_size(self.config, default=1024))
        model.config.use_cache = False
        model = apply_lora_for_causal_lm(model, lora_cfg, distributed=int(world_size) > 1)
        model.config.use_cache = False

        logging.info("Building Qwen X-VARS SFT datasets | rank=%s", rank)
        train_sft = VQAXVarsVideoChatGPTSFTDataset(
            train_data,
            config=self.config,
            tokenizer=tokenizer,
            prompt_cfg=prompt_cfg,
            sft_cfg=sft_cfg,
            xvars_cfg={},
        )
        valid_sft = (
            VQAXVarsVideoChatGPTSFTDataset(
                valid_data,
                config=self.config,
                tokenizer=tokenizer,
                prompt_cfg=prompt_cfg,
                sft_cfg=sft_cfg,
                xvars_cfg={},
            )
            if valid_data is not None
            else None
        )
        if len(train_sft) == 0:
            raise ValueError("No multimodal Qwen training rows were usable.")
        metadata["num_train_samples"] = len(train_sft)
        metadata["num_valid_samples"] = len(valid_sft) if valid_sft is not None else 0
        feature_shape = tuple(train_sft[0]["video_spatio_temporal_features"].shape) if len(train_sft) else None
        logging.info(
            "Built Qwen X-VARS SFT datasets | rank=%s | train=%s | valid=%s | feature_shape=%s",
            rank,
            metadata["num_train_samples"],
            metadata["num_valid_samples"],
            feature_shape,
        )
        train_bs, eval_bs = _resolve_sft_per_device_batch_sizes(self.config, sft_cfg)

        optimizer_cfg = get_train_optimizer(self.config)
        execution_cfg = get_train_execution(self.config)
        fp16, bf16 = _resolve_training_precision_flags(self.config, sft_cfg)
        eval_strategy_key = "evaluation_strategy"
        try:
            if "eval_strategy" in inspect.signature(TrainingArguments.__init__).parameters:
                eval_strategy_key = "eval_strategy"
        except Exception:
            eval_strategy_key = "evaluation_strategy"

        training_kwargs = {
            "output_dir": output_dir,
            "per_device_train_batch_size": train_bs,
            "per_device_eval_batch_size": eval_bs,
            "gradient_accumulation_steps": int(execution_cfg.get("acc_grad_iter", sft_cfg.get("gradient_accumulation_steps", 1))),
            "num_train_epochs": float(get_train_epochs(self.config) or sft_cfg.get("num_train_epochs", 1)),
            "learning_rate": float(optimizer_cfg.get("lr", sft_cfg.get("learning_rate", 2e-4))),
            "logging_steps": int(execution_cfg.get("log_interval", sft_cfg.get("logging_steps", 1))),
            "optim": str(optimizer_cfg.get("hf_optim", "paged_adamw_8bit")),
            "weight_decay": float(optimizer_cfg.get("weight_decay", 0.001)),
            "dataloader_num_workers": int(sft_cfg.get("dataloader_num_workers", 0)),
            "seed": get_system_seed(self.config),
            "data_seed": get_system_seed(self.config),
            "lr_scheduler_type": str(sft_cfg.get("lr_scheduler_type", "constant")),
            "save_strategy": str(sft_cfg.get("save_strategy", "epoch")),
            "report_to": ["wandb"] if use_wandb else [],
            "remove_unused_columns": False,
            "disable_tqdm": bool(sft_cfg.get("disable_tqdm", True)),
            "use_cpu": not get_hf_prefer_cuda(self.config, hf_cfg),
            "fp16": fp16,
            "bf16": bf16,
            "gradient_checkpointing": bool(sft_cfg.get("gradient_checkpointing", False)),
        }
        max_steps = sft_cfg.get("max_steps")
        if max_steps is not None and int(max_steps) > 0:
            training_kwargs["max_steps"] = int(max_steps)
        if int(world_size) > 1:
            training_kwargs["ddp_find_unused_parameters"] = bool(sft_cfg.get("ddp_find_unused_parameters", False))
            training_kwargs["gradient_checkpointing_kwargs"] = {"use_reentrant": False}
        training_kwargs[eval_strategy_key] = _normalize_eval_strategy(
            sft_cfg.get("evaluation_strategy", "epoch"),
            default="epoch",
        )
        args = TrainingArguments(**training_kwargs)
        generated_callback = None
        if bool(generated_validation_cfg.get("enabled", False)):
            generated_callback = _build_xvars_generated_validation_callback(
                tokenizer=tokenizer,
                rows=train_sft.rows,
                output_dir=output_dir,
                validation_cfg=generated_validation_cfg,
                use_step_schedule=bool(training_kwargs.get("max_steps", 0)),
            )
        trainer = XVarsVideoChatGPTTrainer(
            model=model,
            tokenizer=tokenizer,
            args=args,
            train_dataset=train_sft,
            eval_dataset=valid_sft,
            data_collator=XVarsVideoChatGPTDataCollator(tokenizer),
            callbacks=[generated_callback] if generated_callback is not None else None,
        )
        logging.info(
            "Starting Qwen trainer.train | rank=%s | world_size=%s | execution_mode=%s | precision=%s | "
            "max_seq_length=%s | lora_targets=%s | gradient_checkpointing=%s | "
            "ddp_find_unused_parameters=%s | use_cache=%s",
            rank,
            world_size,
            "ddp" if int(world_size) > 1 else "single_gpu",
            "fp16" if fp16 else "bf16" if bf16 else "fp32",
            int(sft_cfg.get("max_seq_length", 768)),
            list(metadata["lora_target_modules"]),
            training_kwargs["gradient_checkpointing"],
            training_kwargs.get("ddp_find_unused_parameters"),
            model.config.use_cache,
        )
        trainer.train()
        logging.info("Finished Qwen trainer.train | rank=%s", rank)
        if rank == 0 and bool(checkpoint_cfg.get("save_adapter", True)):
            trainer.model.save_pretrained(output_dir)
            tokenizer.save_pretrained(output_dir)
            _save_qwen_mm_projector(output_dir, trainer.model)
        metadata["status"] = "trained"
        selected_output_dir = output_dir
        if rank == 0 and generated_callback is not None:
            generated_path = os.path.join(output_dir, "generated_validation.json")
            with open(generated_path, "w", encoding="utf-8") as f:
                json.dump(generated_callback.history, f, indent=2)
            accepted = any(record.get("accepted") for record in generated_callback.history)
            metadata["generated_validation_accepted"] = accepted
            metadata["generated_validation_history"] = generated_path
            if generated_callback.history:
                selected_output_dir = os.path.join(output_dir, "generated_validation_best")
                metadata["best_generated_checkpoint"] = selected_output_dir
                if os.path.isdir(selected_output_dir):
                    _save_qwen_mm_projector(selected_output_dir, trainer.model)
        if rank == 0:
            self._write_artifacts(output_dir, metadata)
            if selected_output_dir != output_dir:
                self._write_artifacts(selected_output_dir, metadata)
            if bool(generated_validation_cfg.get("require_relevance", False)) and not metadata.get(
                "generated_validation_accepted", False
            ):
                raise RuntimeError(
                    "Qwen generated-answer relevance gate failed. Inspect generated_validation.json before "
                    "starting or accepting a full training run."
                )
            return selected_output_dir
        return output_dir

    def _write_artifacts(self, output_dir: str, metadata: dict[str, Any]) -> str:
        return _save_training_artifacts(
            self.config,
            output_dir,
            metadata,
            marker_text="Qwen X-VARS LoRA adapter artifacts are stored in the parent directory when training runs.\n",
        )

VQAQwenVLNativeLoraTrainer

Native Qwen VL LoRA trainer using end-to-end multimodal inputs.

Source code in opensportslib/core/trainer/vqa_trainer.py

class VQAQwenVLNativeLoraTrainer:
    """Native Qwen VL LoRA trainer using end-to-end multimodal inputs."""

    def __init__(self, config):
        self.config = config

    def train(
        self,
        train_data,
        valid_data=None,
        *,
        rank: int = 0,
        world_size: int = 1,
        resume_from_checkpoint: str | None = None,
        use_wandb: bool = False,
    ) -> str:
        execution = get_train_execution(self.config)
        prompt_cfg = _as_dict(execution.get("prompt"))
        sft_cfg = _as_dict(execution.get("sft"))
        hf_cfg = _as_dict(execution.get("hf"))
        lora_cfg = _as_dict(execution.get("lora"))
        checkpoint_cfg = _as_dict(execution.get("checkpoint"))

        model_id = get_vqa_decoder_model_id(self.config, default="Qwen/Qwen2.5-VL-7B-Instruct")
        if "awq" in str(model_id).lower():
            raise ValueError("AWQ Qwen VL checkpoints are inference-only. Use a non-AWQ Qwen VL model for training.")

        save_root = get_system_path(self.config, "save_dir", "./checkpoints") or "./checkpoints"
        output_dir = os.path.join(save_root, "qwen_vl_native_lora")
        os.makedirs(output_dir, exist_ok=True)

        metadata = {
            "backend": "qwen_vl_native_lora",
            "model_id": model_id,
            "status": "metadata_only",
            "multimodal_training": True,
            "visual_config": get_vqa_native_visual_cfg(self.config),
            "num_train_samples": 0,
            "num_valid_samples": 0,
            "invalid_train_rows": 0,
            "invalid_valid_rows": 0,
            "invalid_train_rows_report_path": None,
            "invalid_valid_rows_report_path": None,
            "native_memory_cache_rows": 0,
        }
        if bool(execution.get("dry_run", False)):
            metadata["num_train_samples"] = sum(
                1 for x in train_data if x.get("video_path") or x.get("frame_paths") or x.get("video_frames")
            )
            if valid_data is not None:
                metadata["num_valid_samples"] = sum(
                    1 for x in valid_data if x.get("video_path") or x.get("frame_paths") or x.get("video_frames")
                )
            return self._write_artifacts(output_dir, metadata)

        require_optional_package("transformers", "pip install transformers")
        require_optional_package("peft", "pip install peft")
        from transformers import TrainingArguments

        native_model = QwenVLNativeModel(self.config, model_id=model_id)
        if not native_model._ready:
            raise RuntimeError(native_model._error or "Qwen VL native backend failed to initialize for training.")
        model = native_model.model
        model.config.use_cache = False
        model = apply_lora_for_causal_lm(model, lora_cfg, distributed=int(world_size) > 1)
        model.config.use_cache = False
        cache_tokenized_rows = bool(sft_cfg.get("cache_tokenized_rows", False))
        native_cache_root = os.path.join(output_dir, "native_sft_cache") if cache_tokenized_rows else ""
        train_cache_dir = os.path.join(native_cache_root, "train") if cache_tokenized_rows else None
        valid_cache_dir = os.path.join(native_cache_root, "valid") if cache_tokenized_rows else None
        invalid_train_report_path = os.path.join(output_dir, "native_invalid_rows_train.json")
        invalid_valid_report_path = os.path.join(output_dir, "native_invalid_rows_valid.json")
        metadata["native_sft_cache_enabled"] = cache_tokenized_rows
        metadata["native_sft_cache_dir"] = native_cache_root or None

        train_sft = VQANativeQwenVLSFTDataset(
            train_data,
            model=native_model,
            prompt_cfg=prompt_cfg,
            sft_cfg=sft_cfg,
            cache_dir=train_cache_dir,
            split_name="train",
            invalid_row_report_path=invalid_train_report_path,
            fail_on_invalid=int(world_size) > 1,
            rank=rank,
        )
        valid_sft = (
            VQANativeQwenVLSFTDataset(
                valid_data,
                model=native_model,
                prompt_cfg=prompt_cfg,
                sft_cfg=sft_cfg,
                cache_dir=valid_cache_dir,
                split_name="valid",
                invalid_row_report_path=invalid_valid_report_path,
                fail_on_invalid=int(world_size) > 1,
                rank=rank,
            )
            if valid_data is not None
            else None
        )
        if len(train_sft) == 0:
            raise ValueError("No multimodal native Qwen VL training rows were usable.")
        metadata["num_train_samples"] = len(train_sft)
        metadata["num_valid_samples"] = len(valid_sft) if valid_sft is not None else 0
        metadata["native_memory_cache_rows"] = train_sft.memory_cache_rows
        train_bs, eval_bs = _resolve_sft_per_device_batch_sizes(self.config, sft_cfg)
        optimizer_cfg = get_train_optimizer(self.config)
        execution_cfg = get_train_execution(self.config)
        fp16, bf16 = _resolve_training_precision_flags(self.config, sft_cfg)
        effective_optim = str(
            optimizer_cfg.get("hf_optim")
            or optimizer_cfg.get("type")
            or "paged_adamw_8bit"
        ).strip()
        if effective_optim.lower() == "adamw":
            effective_optim = "adamw_torch"
        eval_strategy_key = "evaluation_strategy"
        try:
            if "eval_strategy" in inspect.signature(TrainingArguments.__init__).parameters:
                eval_strategy_key = "eval_strategy"
        except Exception:
            eval_strategy_key = "evaluation_strategy"

        training_kwargs = {
            "output_dir": output_dir,
            "per_device_train_batch_size": train_bs,
            "per_device_eval_batch_size": eval_bs,
            "gradient_accumulation_steps": int(execution_cfg.get("acc_grad_iter", sft_cfg.get("gradient_accumulation_steps", 1))),
            "num_train_epochs": float(get_train_epochs(self.config) or sft_cfg.get("num_train_epochs", 1)),
            "learning_rate": float(optimizer_cfg.get("lr", sft_cfg.get("learning_rate", 2e-4))),
            "logging_steps": int(execution_cfg.get("log_interval", sft_cfg.get("logging_steps", 1))),
            "optim": effective_optim,
            "weight_decay": float(optimizer_cfg.get("weight_decay", 0.001)),
            "dataloader_num_workers": int(sft_cfg.get("dataloader_num_workers", 0)),
            "seed": get_system_seed(self.config),
            "data_seed": get_system_seed(self.config),
            "lr_scheduler_type": str(sft_cfg.get("lr_scheduler_type", "constant")),
            "save_strategy": str(sft_cfg.get("save_strategy", "epoch")),
            "report_to": ["wandb"] if use_wandb else [],
            "remove_unused_columns": False,
            "disable_tqdm": bool(sft_cfg.get("disable_tqdm", True)),
            "use_cpu": not get_hf_prefer_cuda(self.config, hf_cfg),
            "fp16": fp16,
            "bf16": bf16,
            "gradient_checkpointing": bool(sft_cfg.get("gradient_checkpointing", False)),
        }
        max_steps = sft_cfg.get("max_steps")
        if max_steps is not None and int(max_steps) > 0:
            training_kwargs["max_steps"] = int(max_steps)
        max_grad_norm = sft_cfg.get("max_grad_norm")
        if max_grad_norm is not None:
            training_kwargs["max_grad_norm"] = float(max_grad_norm)
        save_steps = sft_cfg.get("save_steps")
        if save_steps is not None and int(save_steps) > 0:
            training_kwargs["save_steps"] = int(save_steps)
        save_total_limit = sft_cfg.get("save_total_limit")
        if save_total_limit is not None and int(save_total_limit) > 0:
            training_kwargs["save_total_limit"] = int(save_total_limit)
        if int(world_size) > 1:
            training_kwargs["ddp_find_unused_parameters"] = bool(sft_cfg.get("ddp_find_unused_parameters", False))
            ddp_broadcast_buffers = sft_cfg.get("ddp_broadcast_buffers")
            if ddp_broadcast_buffers is not None:
                training_kwargs["ddp_broadcast_buffers"] = bool(ddp_broadcast_buffers)
            training_kwargs["gradient_checkpointing_kwargs"] = {"use_reentrant": False}
        training_kwargs[eval_strategy_key] = _normalize_eval_strategy(
            sft_cfg.get("evaluation_strategy", "epoch"),
            default="epoch",
        )
        logging.info(
            "Starting native Qwen VL trainer.train | rank=%s | world_size=%s | optim=%s | learning_rate=%s | "
            "weight_decay=%s | memory_cache_rows=%s | disk_cache=%s",
            rank,
            world_size,
            effective_optim,
            training_kwargs["learning_rate"],
            training_kwargs["weight_decay"],
            train_sft.memory_cache_rows,
            bool(cache_tokenized_rows),
        )
        args = TrainingArguments(**training_kwargs)
        trainer = NativeQwenVLTrainer(
            model=model,
            args=args,
            train_dataset=train_sft,
            eval_dataset=valid_sft,
            data_collator=NativeQwenVLDataCollator(),
        )
        resume_optimizer_state = bool(sft_cfg.get("resume_optimizer_state", False))
        effective_resume_checkpoint = _prepare_hf_resume_checkpoint(
            resume_from_checkpoint,
            args,
            output_dir,
            resume_optimizer_state=resume_optimizer_state,
        )
        if rank == 0 and resume_from_checkpoint:
            logging.info(
                "Resuming native Qwen VL training from checkpoint: %s | effective_checkpoint=%s | "
                "resume_optimizer_state=%s",
                resume_from_checkpoint,
                effective_resume_checkpoint,
                resume_optimizer_state,
            )
        trainer.train(resume_from_checkpoint=effective_resume_checkpoint)
        metadata["invalid_train_rows"] = train_sft.invalid_row_count
        metadata["invalid_valid_rows"] = valid_sft.invalid_row_count if valid_sft is not None else 0
        metadata["invalid_train_rows_report_path"] = (
            invalid_train_report_path if train_sft.invalid_row_count > 0 else None
        )
        metadata["invalid_valid_rows_report_path"] = (
            invalid_valid_report_path if valid_sft is not None and valid_sft.invalid_row_count > 0 else None
        )
        if rank == 0 and bool(checkpoint_cfg.get("save_adapter", True)):
            trainer.model.save_pretrained(output_dir)
            if getattr(native_model, "processor", None) is not None:
                native_model.processor.save_pretrained(output_dir)
        metadata["status"] = "trained"
        if rank == 0:
            self._write_artifacts(output_dir, metadata)
        return output_dir

    def _write_artifacts(self, output_dir: str, metadata: dict[str, Any]) -> str:
        return _save_training_artifacts(
            self.config,
            output_dir,
            metadata,
            marker_text="Qwen VL native LoRA adapter artifacts are stored in the parent directory when training runs.\n",
        )

Trainer_VQA

VQA trainer dispatcher plus inference/evaluation helpers.

Source code in opensportslib/core/trainer/vqa_trainer.py

class Trainer_VQA:
    """VQA trainer dispatcher plus inference/evaluation helpers."""

    def __init__(self, config):
        self.config = config
        self.best_checkpoint_path: str | None = None
        self.loaded_checkpoint_metadata: dict[str, Any] | None = None

    def load(self, weights: str):
        # VQA LoRA checkpoints are directory-based adapter artifacts.
        metadata_path = os.path.join(weights, "training_metadata.json") if os.path.isdir(weights) else ""
        if metadata_path and os.path.exists(metadata_path):
            with open(metadata_path, encoding="utf-8") as f:
                self.loaded_checkpoint_metadata = json.load(f)
            if self.loaded_checkpoint_metadata.get("generated_validation_accepted") is False:
                raise ValueError(
                    "Refusing to load a rejected VQA adapter: generated validation did not pass. "
                    f"Checkpoint: {weights}. Use the configured base model or a validated adapter."
                )
        self.best_checkpoint_path = weights
        return weights

    def train(
        self,
        model,
        train_data,
        valid_data=None,
        *,
        rank: int = 0,
        world_size: int = 1,
        resume_from_checkpoint: str | None = None,
        use_wandb: bool = False,
    ) -> str:
        del model
        execution = get_train_execution(self.config)
        backend = str(execution.get("training_backend", "placeholder")).lower()
        if backend == "xvars_videochatgpt_lora":
            ckpt = VQAXVarsVideoChatGPTLoraTrainer(self.config).train(
                train_data,
                valid_data,
                rank=rank,
                world_size=world_size,
                use_wandb=use_wandb,
            )
            self.best_checkpoint_path = ckpt
            return ckpt
        if backend == "qwen_xvars_lora":
            ckpt = VQAQwenXVarsLoraTrainer(self.config).train(
                train_data,
                valid_data,
                rank=rank,
                world_size=world_size,
                use_wandb=use_wandb,
            )
            self.best_checkpoint_path = ckpt
            return ckpt
        if backend == "qwen_vl_native_lora":
            ckpt = VQAQwenVLNativeLoraTrainer(self.config).train(
                train_data,
                valid_data,
                rank=rank,
                world_size=world_size,
                resume_from_checkpoint=resume_from_checkpoint,
                use_wandb=use_wandb,
            )
            self.best_checkpoint_path = ckpt
            return ckpt
        raise ValueError(
            f"Unsupported VQA training backend '{backend}'. "
            "Only 'xvars_videochatgpt_lora', 'qwen_xvars_lora', and 'qwen_vl_native_lora' are supported."
        )

    def infer(self, model, dataset, *, use_wandb: bool = False) -> dict[str, Any]:
        exec_cfg = get_train_execution(self.config)
        prompt_cfg = exec_cfg.get("prompt", {}) if isinstance(exec_cfg, dict) else {}
        generation_cfg = get_vqa_generation_cfg(self.config)

        preds = []
        disable_tqdm = bool(exec_cfg.get("disable_tqdm", False) if isinstance(exec_cfg, dict) else False)
        try:
            total = len(dataset)
        except Exception:
            total = None
        logging.info(
            "Starting VQA inference | samples=%s | max_new_tokens=%s",
            total if total is not None else "unknown",
            generation_cfg.get("max_new_tokens"),
        )
        started_at = time.perf_counter()
        iterator = tqdm(
            dataset,
            total=total,
            desc="VQA inference",
            unit="sample",
            disable=disable_tqdm,
            leave=False,
        )
        for sample in iterator:
            sample_started_at = time.perf_counter()
            answer = model.generate_answer(
                sample,
                prompt_cfg=prompt_cfg,
                generation_cfg=generation_cfg,
            )
            if not disable_tqdm:
                iterator.set_postfix_str(f"id={sample.get('id')} elapsed={time.perf_counter() - sample_started_at:.2f}s")
            row = {
                "id": sample.get("id"),
                "question": sample.get("question"),
                "answer_text": answer,
                "video_path": sample.get("video_path"),
            }
            row.update(
                _normalize_classification_answer(
                    answer,
                    sample.get("allowed_labels"),
                    sample.get("ground_truth_label"),
                )
            )
            preds.append(row)
        logging.info(
            "Finished VQA inference | samples=%s | elapsed_s=%.2f",
            len(preds),
            time.perf_counter() - started_at,
        )
        payload = {"task": "vqa", "data": preds}
        _maybe_log_vqa_predictions(payload, use_wandb=use_wandb)
        return payload

    def evaluate(self, predictions: dict[str, Any], dataset, *, use_wandb: bool = False) -> dict[str, Any]:
        metrics = compute_vqa_metrics(predictions, dataset, eval_profile=get_vqa_eval_profile_cfg(self.config))
        _maybe_log_vqa_metrics(metrics, use_wandb=use_wandb)
        return metrics

build_vqa_sft_text(sample, *, config=None, prompt_cfg=None, sft_cfg=None, reference=None)

Convert a VQADataset sample into prompt/answer/text fields for SFT.

Source code in opensportslib/core/trainer/vqa_trainer.py

def build_vqa_sft_text(
    sample: dict[str, Any],
    *,
    config=None,
    prompt_cfg: dict[str, Any] | None = None,
    sft_cfg: dict[str, Any] | None = None,
    reference: str | None = None,
) -> dict[str, str]:
    """Convert a VQADataset sample into prompt/answer/text fields for SFT."""
    prompt_cfg = prompt_cfg or {}
    sft_cfg = sft_cfg or {}
    refs = sample.get("references") or []
    answer = str(reference).strip() if reference is not None else (str(refs[0]).strip() if refs else "")
    question = str(sample.get("question", "")).strip()
    system_prompt = str(
        prompt_cfg.get(
            "system_prompt",
            "You are an artificial intelligence assistant for visual question answering.",
        )
    ).strip()

    prior_text = ""
    if bool(prompt_cfg.get("include_priors", True)):
        built_prior = build_prior_text(
            sample.get("labels", {}) or {},
            sample.get("metadata", {}) or {},
            include_fields=prompt_cfg.get("prior_fields"),
        )
        prior_text = str(sample.get("prior_prediction_text", "")).strip() or built_prior

    if config is not None:
        token_len = _resolve_vqa_video_token_len(config, prompt_cfg=prompt_cfg, sft_cfg=sft_cfg)
    else:
        token_len = int(prompt_cfg.get("video_token_len", sft_cfg.get("video_token_len", 300)))
    token_len = token_len if bool(sft_cfg.get("include_video_tokens", True)) else 0
    prompt = build_xvars_prompt(
        system_prompt=system_prompt,
        question=question,
        prior_text=prior_text,
        video_token_len=token_len,
    )
    append_eos = bool(sft_cfg.get("append_eos_token", True))
    completion = f"{answer}</s>" if answer and append_eos else answer
    return {
        "prompt": prompt,
        "answer": answer,
        "completion": completion,
        "text": f"{prompt} {completion}".strip(),
    }

VQA runtime

The VQA module keeps backend-specific collators, supervised fine-tuning datasets, and LoRA trainer classes beside Trainer_VQA. They are internal implementation classes selected by the configured backend and training mode; they may save adapters, checkpoints, and metadata in the configured output directory. Use VQAModel.train(), infer(), and evaluate() for the public contract.

Metric signatures

opensportslib/metrics/ is a namespace-style source directory without an __init__.py, so MkDocstrings cannot resolve it as a dotted package. Its non-private metric entry points are maintained as source-linked signatures:

# classification_metric.py
process_preds_labels(eval_pred, top_k=None)
compute_classification_metrics(eval_pred, top_k=None, mode="logits")
compute_detailed_classification_metrics(all_logits, all_labels, class_names, save_dir, set_name)

# localization_metric.py
compute_average_precision(pred, truth, tolerance=0, min_precision=0, plot_ax=None, plot_label=None, plot_raw_pr=True)
compute_mAPs_E2E(truth, pred, tolerances=[0, 1, 2, 3, 4], plot_pr=False)
build_snpro_prediction_json(pred_events, head_name="action", split=None, created_by="model")
infer_and_process_predictions_e2e(...)
compute_performances_mAP(...)

# vqa_metric.py
compute_vqa_metrics(predictions, dataset, eval_profile=None)

These functions are extension points. Their source files are metrics/classification_metric.py, metrics/localization_metric.py, and metrics/vqa_metric.py; detailed localization arguments should be read from the source before extending an evaluation route.