Function evaluate

tools/test_linear_probing_hf.py:303–345

Evaluate all classifiers and return accuracies.

(
    feature_model: nn.Module,
    linear_classifiers: AllClassifiers,
    val_loader: DataLoader,
    device: torch.device,
)

@torch.no_grad()
def evaluate(
    feature_model: nn.Module,
    linear_classifiers: AllClassifiers,
    val_loader: DataLoader,
    device: torch.device,
) -> Dict[str, float]:
    """Evaluate all classifiers and return accuracies."""
    linear_classifiers.eval()

    classifiers_dict = linear_classifiers.module.classifiers_dict if hasattr(linear_classifiers, 'module') else linear_classifiers.classifiers_dict

    correct = {k: 0 for k in classifiers_dict.keys()}
    total = 0

    progress_bar = tqdm(val_loader, desc="Evaluating") if is_main_process() else val_loader

    for images, labels in progress_bar:
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        features = feature_model(images)
        outputs = linear_classifiers(features)

        for k, logits in outputs.items():
            preds = logits.argmax(dim=1)
            correct[k] += (preds == labels).sum().item()

        total += labels.size(0)

    # Aggregate across processes
    if dist.is_initialized():
        # Gather correct counts and total from all processes
        total_tensor = torch.tensor([total], device=device, dtype=torch.long)
        dist.all_reduce(total_tensor)
        total = total_tensor.item()

        for k in correct:
            correct_tensor = torch.tensor([correct[k]], device=device, dtype=torch.long)
            dist.all_reduce(correct_tensor)
            correct[k] = correct_tensor.item()

    accuracies = {k: 100.0 * v / total for k, v in correct.items()}
    return accuracies


# ============================================================================