FixMatch

def set_data_loader(
    self,
    cfg: DotMap,
    lb_dset: BasicDataset,
    ulb_dset: BasicDataset,
    eval_dset: BasicDataset | None = None,
) -> None:
    """Set up data loaders for training and evaluation.

    Args:
        cfg: Configuration object with dataloader parameters
        lb_dset: Labeled dataset
        ulb_dset: Unlabeled dataset
        eval_dset: Evaluation dataset (optional)
    """
    logger.debug("Setting up data loaders")

    loader_dict = {}
    dset_dict = {"train_lb": lb_dset, "train_ulb": ulb_dset, "eval": eval_dset}

    loader_dict["train_lb"] = get_data_loader(
        dset_dict["train_lb"],
        cfg.batch_size,
        data_sampler="RandomSampler",
        num_iters=cfg.num_train_iter,
        num_workers=cfg.num_workers,
        pin_memory=cfg.pin_memory,
        use_weighted_sampler=cfg.oversample,
    )

    loader_dict["train_ulb"] = get_data_loader(
        dset_dict["train_ulb"],
        cfg.batch_size * cfg.uratio,
        data_sampler="RandomSampler",
        num_iters=cfg.num_train_iter,
        num_workers=cfg.num_workers,
        pin_memory=cfg.pin_memory,
    )

    # Only create eval loader if eval_dset is provided
    if eval_dset is not None:
        actual_eval_batch_size = min(cfg.eval_batch_size, len(dset_dict["eval"]))
        loader_dict["eval"] = get_data_loader(
            dset_dict["eval"], actual_eval_batch_size, num_workers=0
        )

    self.loader_dict = loader_dict

def set_optimizer(
    self,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler | None = None,
) -> None:
    """Set the optimizer and optional scheduler.

    Args:
        optimizer: PyTorch optimizer for training
        scheduler: Optional learning rate scheduler
    """
    self.optimizer = optimizer
    self.scheduler = scheduler

def train(
    self, cfg: DotMap, progressbar: Any = None, progress_callback: Callable | None = None
) -> dict[str, Any]:
    """Train the model using the FixMatch algorithm.

    Args:
        cfg: Configuration object with training parameters
        progressbar: Optional tqdm progress bar
        progress_callback: Optional callback function for reporting progress

    Returns:
        dict: Evaluation metrics after training
    """
    ngpus_per_node = torch.cuda.device_count()

    logger.info(
        f"Starting FixMatch training for {cfg.num_train_iter} iterations on {ngpus_per_node} GPUs"
        + f" with normalisation {cfg.normalisation.normalisation_method.name}"
    )

    self.it = 0

    # Set model to training mode
    self.train_model.train()

    # Set up progress tracking
    progressbar = tqdm(
        desc="Training...",
        total=cfg.num_train_iter + self.total_it,
        file=sys.stdout,
        disable=(cfg.log_level not in ["DEBUG", "TRACE"]),
    )
    progressbar.update(self.total_it)
    progressbar.refresh()

    # Main training loop
    for (x_lb, y_lb, _), (x_ulb_w, x_ulb_s, _) in zip(
        self.loader_dict["train_lb"], self.loader_dict["train_ulb"]
    ):
        # Check if we've reached target iterations
        if self.it >= cfg.num_train_iter:
            logger.debug(f"Training finished after {self.total_it} iterations")
            break

        # Get batch sizes
        num_lb = x_lb.shape[0]
        num_ulb = x_ulb_w.shape[0]
        assert num_ulb == x_ulb_s.shape[0]

        # Log training details at TRACE level
        logger.trace(
            f"Iteration {self.it}: labeled batch shape: {x_lb.shape}, "
            f"label distribution: {y_lb.unique(return_counts=True)}"
        )

        # Move to GPU if available
        if torch.cuda.is_available():
            x_lb, x_ulb_w, x_ulb_s = (
                x_lb.cuda(cfg.gpu),
                x_ulb_w.cuda(cfg.gpu),
                x_ulb_s.cuda(cfg.gpu),
            )
            y_lb = y_lb.cuda(cfg.gpu)

        # Combine inputs for efficient processing
        inputs = torch.cat((x_lb, x_ulb_w, x_ulb_s))

        # Forward pass and loss calculation
        logits = self.train_model(inputs)
        logits_x_lb = logits[:num_lb]
        logits_x_ulb_w, logits_x_ulb_s = logits[num_lb:].chunk(2)
        del logits  # Free memory

        # Calculate supervised loss
        sup_loss = cross_entropy_loss(logits_x_lb, y_lb, reduction="mean")

        # Calculate unsupervised loss with consistency regularization
        unsup_loss, mask = consistency_loss(
            logits_x_ulb_w,
            logits_x_ulb_s,
            "ce",
            self.T,
            self.p_cutoff,
            use_hard_labels=cfg.hard_label,
        )

        # Combine losses
        total_loss = sup_loss + self.lambda_u * unsup_loss

        # Backpropagation and parameter updates
        total_loss.backward()
        self.optimizer.step()

        if self.scheduler:
            self.scheduler.step()
        self.train_model.zero_grad()

        # Update evaluation model with EMA
        with torch.no_grad():
            self._eval_model_update()
            train_accuracy = accuracy(logits_x_lb, y_lb)[0]

        # Store metrics for logging
        metrics = {
            "sup_loss": sup_loss.detach().item(),
            "unsup_loss": unsup_loss.detach().item(),
            "total_loss": total_loss.detach().item(),
            "mask_ratio": (1.0 - mask.detach()).item(),
            "learning_rate": self.optimizer.param_groups[0]["lr"],
            "accuracy": train_accuracy.item(),
        }

        # Update progress information
        progressbar.set_postfix_str(f"Loss={metrics['total_loss']:.3e}")
        progressbar.update(1)
        if progress_callback:
            progress_callback(self.it, cfg.num_train_iter)
        progressbar.refresh()

        # Periodic evaluation
        if cfg.num_eval_iter > 0 and self.it % cfg.num_eval_iter == 0 and self.it > 0:
            progressbar.close()
            logger.debug(f"Performing evaluation at {self.total_it} iterations")

            # Run evaluation
            eval_dict = self.evaluate(cfg=cfg)

            # Update best accuracy if needed
            if eval_dict["eval/top-1-acc"] > self.best_eval_acc:
                self.best_eval_acc = eval_dict["eval/top-1-acc"]
                self.best_it = self.total_it

            # Log evaluation results
            self.print_fn(
                f"Iteration {self.total_it}, EMA model used: {hasattr(self, 'eval_model')}, "
                f"Accuracy: {eval_dict['eval/top-1-acc']:.4f}, Loss: {eval_dict['eval/loss']:.4e}, "
                f"AUROC: {eval_dict['eval/auroc']:.4f}, AUPRC: {eval_dict['eval/auprc']:.4f}, "
                f"Best accuracy: {self.best_eval_acc:.4f} at iteration {self.best_it}"
            )

            # Create new progress bar
            progressbar = tqdm(
                desc="Training...",
                total=cfg.num_train_iter + self.total_it - self.it,
                file=sys.stdout,
            )
            progressbar.update(self.total_it)
            progressbar.refresh()

        # Increment iteration counters
        self.it += 1
        self.total_it += 1

        # Update session tracker with model loss if available
        if self.session_tracker is not None:
            self.session_tracker.update_model_iteration(total_loss.item())

    progressbar.close()

    # Final evaluation after training
    eval_dict = self.evaluate(cfg=cfg)
    return eval_dict

@torch.no_grad()
def evaluate(
    self,
    cfg: DotMap,
    eval_loader: DataLoader | None = None,
    progress_callback: Callable | None = None,
) -> dict[str, Any]:
    """Evaluate the model on a dataset.

    Args:
        cfg: Configuration object
        eval_loader: Optional custom evaluation dataloader
        progress_callback: Optional callback function for reporting progress

    Returns:
        dict: Dictionary of evaluation metrics
    """
    if cfg.test_ratio <= 0:
        logger.info("Test ratio is 0, skipping evaluation")
        return {}

    # Clear GPU memory before evaluation
    torch.cuda.empty_cache()

    # Determine which model to use for evaluation
    use_ema = hasattr(self, "eval_model")
    eval_model = self.eval_model if use_ema else self.train_model
    eval_model.eval()

    # Use provided loader or default
    if eval_loader is None:
        eval_loader = self.loader_dict["eval"]

    # Initialize metrics
    total_loss = 0.0
    total_acc = 0.0
    total_num = 0.0
    total_len = len(eval_loader)
    confusion_matrix = torch.zeros(2, 2)
    predictions_and_labels = {}
    all_probs = []
    all_labels = []

    # Evaluate all batches
    logger.debug(f"Starting evaluation over {total_len} batches")
    for batch_idx, (x, y, filename) in enumerate(eval_loader):
        # Move to GPU if available
        if torch.cuda.is_available():
            x, y = x.cuda(cfg.gpu), y.cuda(cfg.gpu)

        # Process batch
        num_batch = x.shape[0]
        total_num += num_batch

        # Forward pass
        logits = eval_model(x)
        loss = F.cross_entropy(logits, y.long(), reduction="mean")
        pred = torch.max(logits, dim=-1)[1]
        probs = F.softmax(logits, dim=-1)[:, 1]  # Probability of anomaly class

        # Calculate accuracy
        acc = torch.sum(pred == y)

        # Update confusion matrix
        for i in range(y.shape[0]):
            confusion_matrix[y[i], pred[i]] += 1

        # Accumulate metrics
        total_loss += loss.detach() * num_batch
        total_acc += acc.detach()

        # Store data for ROC and PR curves
        all_probs.extend(probs.cpu().numpy())
        all_labels.extend(y.cpu().numpy())

        # Store predictions and labels
        predictions_and_labels.update(
            {filename[i]: (probs[i].detach(), y[i].detach()) for i in range(num_batch)}
        )

        # Update progress
        if progress_callback:
            progress_callback(batch_idx + 1, total_len)

    # Restore train mode if necessary
    if not use_ema:
        eval_model.train()

    # Calculate final metrics
    if len(all_labels) > 0 and len(set(all_labels)) > 1:
        # Compute AUROC and AUPRC
        auroc = roc_auc_score(all_labels, all_probs)
        precision, recall, _ = precision_recall_curve(all_labels, all_probs)
        auprc = auc(recall, precision)

        logger.debug(f"Evaluated {total_num} samples")
        logger.debug(
            f"Loss: {total_loss / total_num:.4f}, Top-1 Acc: {total_acc / total_num:.4f}"
        )
        logger.info(f"AUROC: {auroc:.4f}, AUPRC: {auprc:.4f}")
        logger.info(f"Confusion Matrix: {confusion_matrix}")

        # Sort predictions by confidence difference (most uncertain first)
        predictions_and_labels = {
            k: v
            for k, v in sorted(
                predictions_and_labels.items(),
                key=lambda item: abs(item[1][0].item() - item[1][1].item()),
                reverse=True,
            )
        }
    else:
        logger.warning("Not enough distinct class labels for ROC/PRC calculation")
        auroc = 0.0
        auprc = 0.0
        precision = []
        recall = []

    # Return metrics dictionary
    return {
        "eval/loss": float(total_loss / total_num) if total_num > 0 else float("inf"),
        "eval/top-1-acc": float(total_acc / total_num) if total_num > 0 else 0.0,
        "eval/auroc": float(auroc),
        "eval/auprc": float(auprc),
        "eval/confusion_matrix": confusion_matrix.cpu().numpy().tolist(),
        "eval/predictions_and_labels": predictions_and_labels,
        "eval/roc_data": (all_labels, all_probs),
        "eval/precision_recall": (precision, recall),
    }