birder 0.4.1__py3-none-any.whl → 0.4.4__py3-none-any.whl

birder/eval/classification.py

@@ -0,0 +1,235 @@
+import argparse
+import logging
+from typing import Any
+
+import torch
+from torch.utils.data import DataLoader
+
+import birder
+from birder.common import cli
+from birder.conf import settings
+from birder.data.dataloader.webdataset import make_wds_loader
+from birder.data.datasets.directory import make_image_dataset
+from birder.data.datasets.webdataset import make_wds_dataset
+from birder.data.datasets.webdataset import prepare_wds_args
+from birder.data.datasets.webdataset import wds_args_from_info
+from birder.inference.data_parallel import InferenceDataParallel
+
+logger = logging.getLogger(__name__)
+
+
+# pylint: disable=too-many-branches
+def evaluate(args: argparse.Namespace) -> None:
+    if args.gpu is True:
+        device = torch.device("cuda")
+    elif args.mps is True:
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    if args.parallel is True and torch.cuda.device_count() > 1:
+        logger.info(f"Using {torch.cuda.device_count()} {device} devices")
+    else:
+        if args.gpu_id is not None:
+            torch.cuda.set_device(args.gpu_id)
+
+        logger.info(f"Using device {device}")
+
+    if args.fast_matmul is True or args.amp is True:
+        torch.set_float32_matmul_precision("high")
+
+    model_dtype: torch.dtype = getattr(torch, args.model_dtype)
+    amp_dtype: torch.dtype = getattr(torch, args.amp_dtype)
+    model_list = birder.list_pretrained_models(args.filter)
+    for model_name in model_list:
+        net, (class_to_idx, signature, rgb_stats, *_) = birder.load_pretrained_model(
+            model_name, inference=True, device=device, dtype=model_dtype
+        )
+        if args.channels_last is True:
+            net = net.to(memory_format=torch.channels_last)
+            logger.debug("Using channels-last memory format")
+
+        if args.parallel is True and torch.cuda.device_count() > 1:
+            net = InferenceDataParallel(net, output_device="cpu", compile_replicas=args.compile)
+        elif args.compile is True:
+            net = torch.compile(net)
+
+        if args.size is None:
+            size = birder.get_size_from_signature(signature)
+        else:
+            size = args.size
+
+        transform = birder.classification_transform(size, rgb_stats, args.center_crop, args.simple_crop)
+
+        if args.wds is True:
+            wds_path: str | list[str]
+            if args.wds_info is not None:
+                wds_path, dataset_size = wds_args_from_info(args.wds_info, args.wds_split)
+                if args.wds_size is not None:
+                    dataset_size = args.wds_size
+            else:
+                wds_path, dataset_size = prepare_wds_args(args.data_path[0], args.wds_size, device)
+
+            num_samples = dataset_size
+            dataset = make_wds_dataset(
+                wds_path,
+                dataset_size=dataset_size,
+                shuffle=False,
+                samples_names=True,
+                transform=transform,
+            )
+            inference_loader = make_wds_loader(
+                dataset,
+                args.batch_size,
+                num_workers=args.num_workers,
+                prefetch_factor=None,
+                collate_fn=None,
+                world_size=1,
+                pin_memory=False,
+                exact=True,
+            )
+        else:
+            dataset = make_image_dataset(args.data_path, class_to_idx, transforms=transform)
+            num_samples = len(dataset)
+            inference_loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+
+        with torch.inference_mode():
+            results = birder.evaluate_classification(
+                device,
+                net,
+                inference_loader,
+                class_to_idx,
+                args.tta,
+                args.channels_last,
+                model_dtype,
+                args.amp,
+                amp_dtype,
+                num_samples=num_samples,
+                sparse=args.save_sparse_results,
+            )
+
+        logger.info(f"{model_name}: accuracy={results.accuracy:.4f}")
+        base_output_path = (
+            f"{args.dir}/{model_name}_{len(class_to_idx)}_{size[0]}px_crop{args.center_crop}_{num_samples}"
+        )
+        if args.save_sparse_results is True:
+            results_file_suffix = "_sparse.csv"
+        else:
+            results_file_suffix = ".csv"
+
+        results.save(f"{base_output_path}{results_file_suffix}")
+
+
+def set_parser(subparsers: Any) -> None:
+    subparser = subparsers.add_parser(
+        "classification",
+        allow_abbrev=False,
+        help="evaluate pretrained classification models on a dataset",
+        description="evaluate pretrained classification models on a dataset",
+        epilog=(
+            "Usage examples:\n"
+            "python -m birder.eval classification --filter '*il-all*' --fast-matmul --gpu "
+            "data/validation_il-all_packed\n"
+            "python -m birder.eval classification --amp --compile --gpu --gpu-id 1 data/testing\n"
+            "python -m birder.eval classification --filter '*inat21*' --amp --compile --gpu "
+            "--parallel ~/Datasets/inat2021/val\n"
+            "python -m birder.eval classification --wds --wds-info data/validation_wds/info.json --gpu\n"
+        ),
+        formatter_class=cli.ArgumentHelpFormatter,
+    )
+    subparser.add_argument("--filter", type=str, help="models to evaluate (fnmatch type filter)")
+    subparser.add_argument("--compile", default=False, action="store_true", help="enable compilation")
+    subparser.add_argument(
+        "--channels-last", default=False, action="store_true", help="use channels-last memory format"
+    )
+    subparser.add_argument(
+        "--model-dtype",
+        type=str,
+        choices=["float32", "float16", "bfloat16"],
+        default="float32",
+        help="model dtype to use",
+    )
+    subparser.add_argument(
+        "--amp", default=False, action="store_true", help="use torch.amp.autocast for mixed precision inference"
+    )
+    subparser.add_argument(
+        "--amp-dtype",
+        type=str,
+        choices=["float16", "bfloat16"],
+        default="float16",
+        help="whether to use float16 or bfloat16 for mixed precision",
+    )
+    subparser.add_argument(
+        "--fast-matmul", default=False, action="store_true", help="use fast matrix multiplication (affects precision)"
+    )
+    subparser.add_argument("--tta", default=False, action="store_true", help="test time augmentation (oversampling)")
+    subparser.add_argument(
+        "--size", type=int, nargs="+", metavar=("H", "W"), help="image size for inference (defaults to model signature)"
+    )
+    subparser.add_argument("--batch-size", type=int, default=64, metavar="N", help="the batch size")
+    subparser.add_argument(
+        "-j", "--num-workers", type=int, default=8, metavar="N", help="number of preprocessing workers"
+    )
+    subparser.add_argument("--center-crop", type=float, default=1.0, help="center crop ratio to use during inference")
+    subparser.add_argument(
+        "--simple-crop",
+        default=False,
+        action="store_true",
+        help="use a simple crop that preserves aspect ratio but may trim parts of the image",
+    )
+    subparser.add_argument(
+        "--dir", type=str, default="evaluate", help="place all outputs in a sub-directory (relative to results)"
+    )
+    subparser.add_argument("--gpu", default=False, action="store_true", help="use gpu")
+    subparser.add_argument("--gpu-id", type=int, metavar="ID", help="gpu id to use")
+    subparser.add_argument(
+        "--mps", default=False, action="store_true", help="use mps (Metal Performance Shaders) device"
+    )
+    subparser.add_argument("--parallel", default=False, action="store_true", help="use multiple gpus")
+    subparser.add_argument(
+        "--save-sparse-results",
+        default=False,
+        action="store_true",
+        help="save results object in memory-efficient sparse format (only top-k probabilities)",
+    )
+    subparser.add_argument("--wds", default=False, action="store_true", help="evaluate a webdataset directory")
+    subparser.add_argument("--wds-size", type=int, metavar="N", help="size of the wds dataset")
+    subparser.add_argument("--wds-info", type=str, metavar="FILE", help="wds info file path")
+    subparser.add_argument(
+        "--wds-split", type=str, default="validation", metavar="NAME", help="wds dataset split to load"
+    )
+    subparser.add_argument("data_path", nargs="*", help="data files path (directories and files)")
+    subparser.set_defaults(func=main)
+
+
+def validate_args(args: argparse.Namespace) -> None:
+    if args.amp is True and args.model_dtype != "float32":
+        raise cli.ValidationError("--amp can only be used with --model-dtype float32")
+    if args.center_crop > 1 or args.center_crop <= 0.0:
+        raise cli.ValidationError(f"--center-crop must be in range of (0, 1.0], got {args.center_crop}")
+    if args.parallel is True and args.gpu is False:
+        raise cli.ValidationError("--parallel requires --gpu to be set")
+    if args.wds is False and len(args.data_path) == 0:
+        raise cli.ValidationError("Must provide at least one data source, --data-path or --wds")
+
+    if args.wds is True:
+        if args.wds_info is None and len(args.data_path) == 0:
+            raise cli.ValidationError("--wds requires a data path unless --wds-info is provided")
+        if len(args.data_path) > 1:
+            raise cli.ValidationError(f"--wds can have at most 1 --data-path, got {len(args.data_path)}")
+        if args.wds_info is None and len(args.data_path) == 1:
+            data_path = args.data_path[0]
+            if "://" in data_path and args.wds_size is None:
+                raise cli.ValidationError("--wds-size is required for remote --data-path")
+
+    args.size = cli.parse_size(args.size)
+
+
+def main(args: argparse.Namespace) -> None:
+    validate_args(args)
+
+    if settings.RESULTS_DIR.joinpath(args.dir).exists() is False:
+        logger.info(f"Creating {settings.RESULTS_DIR.joinpath(args.dir)} directory...")
+        settings.RESULTS_DIR.joinpath(args.dir).mkdir(parents=True)
+
+    evaluate(args)

birder/eval/methods/ami.py

@@ -0,0 +1,78 @@
+"""
+AMI clustering for unsupervised embedding evaluation
+
+Paper "An Empirical Study into Clustering of Unseen Datasets with Self-Supervised Encoders",
+https://arxiv.org/abs/2406.02465
+"""
+
+import logging
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+from sklearn.cluster import AgglomerativeClustering
+from sklearn.metrics import adjusted_mutual_info_score
+
+from birder.eval._embeddings import l2_normalize
+
+logger = logging.getLogger(__name__)
+
+try:
+    import umap
+
+    _HAS_UMAP = True
+except ImportError:
+    _HAS_UMAP = False
+
+
+def evaluate_ami(
+    features: npt.NDArray[np.float32],
+    labels: npt.NDArray[np.int_],
+    n_clusters: int,
+    umap_dim: int = 50,
+    l2_normalize_features: bool = True,
+    seed: Optional[int] = None,
+) -> float:
+    """
+    Evaluate embedding quality using UMAP + Agglomerative Clustering + AMI
+
+    Parameters
+    ----------
+    features
+        Feature array of shape (n_samples, embedding_dim).
+    labels
+        True labels of shape (n_samples,).
+    n_clusters
+        Number of clusters (should match number of true classes).
+    umap_dim
+        Target dimensionality for UMAP reduction.
+    l2_normalize_features
+        If True, applies row-wise L2 normalization before UMAP.
+    seed
+        Random seed for UMAP reproducibility. When None, uses all available
+        cores (n_jobs=-1) but results are non-deterministic. When set, forces n_jobs=1.
+
+    Returns
+    -------
+    ami_score
+        Adjusted Mutual Information score between true labels and cluster assignments.
+    """
+
+    assert _HAS_UMAP, "'pip install umap-learn' to use AMI evaluation"
+
+    if seed is None:
+        logger.warning("No seed set, UMAP results will be non-deterministic (using n_jobs=-1)")
+        n_jobs = -1
+    else:
+        n_jobs = 1
+
+    if l2_normalize_features is True:
+        features = l2_normalize(features)
+
+    reducer = umap.UMAP(n_components=umap_dim, min_dist=0.0, n_jobs=n_jobs, random_state=seed)
+    features_reduced = reducer.fit_transform(features)
+
+    clustering = AgglomerativeClustering(n_clusters=n_clusters, linkage="ward")
+    cluster_assignments = clustering.fit_predict(features_reduced)
+
+    return float(adjusted_mutual_info_score(labels, cluster_assignments))

birder/eval/methods/knn.py

@@ -0,0 +1,71 @@
+"""
+K-Nearest Neighbors classifier for few-shot learning evaluation
+
+Uses cosine similarity (dot product of L2-normalized features) with temperature-scaled softmax voting.
+"""
+
+import numpy as np
+import numpy.typing as npt
+
+from birder.eval._embeddings import l2_normalize
+
+
+def evaluate_knn(
+    train_features: npt.NDArray[np.float32],
+    train_labels: npt.NDArray[np.int_],
+    test_features: npt.NDArray[np.float32],
+    test_labels: npt.NDArray[np.int_],
+    k: int,
+    temperature: float = 0.07,
+) -> tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]:
+    """
+    Evaluate using K-Nearest Neighbors with cosine similarity and soft voting
+
+    Parameters
+    ----------
+    train_features
+        Training features of shape (n_train, embedding_dim).
+    train_labels
+        Training labels of shape (n_train,).
+    test_features
+        Test features of shape (n_test, embedding_dim).
+    test_labels
+        Test labels of shape (n_test,).
+    k
+        Number of nearest neighbors.
+    temperature
+        Temperature for softmax scaling.
+
+    Returns
+    -------
+    y_pred
+        Predicted labels for test samples.
+    y_true
+        True labels for test samples (same as test_labels).
+    """
+
+    # Cosine similarity
+    train_norm = l2_normalize(train_features)
+    test_norm = l2_normalize(test_features)
+    similarities = test_norm @ train_norm.T
+
+    # Get top-k neighbors
+    top_k_indices = np.argsort(-similarities, axis=1)[:, :k]
+    top_k_sims = np.take_along_axis(similarities, top_k_indices, axis=1)
+    top_k_labels = train_labels[top_k_indices]
+
+    # Temperature-scaled softmax voting
+    top_k_sims_scaled = top_k_sims / temperature
+    top_k_sims_scaled = top_k_sims_scaled - top_k_sims_scaled.max(axis=1, keepdims=True)
+    weights = np.exp(top_k_sims_scaled)
+    weights = weights / weights.sum(axis=1, keepdims=True)
+
+    # Weighted voting
+    num_classes = train_labels.max() + 1
+    votes = np.zeros((len(test_features), num_classes), dtype=np.float32)
+    for i in range(k):
+        np.add.at(votes, (np.arange(len(test_features)), top_k_labels[:, i]), weights[:, i])
+
+    y_pred = votes.argmax(axis=1).astype(np.int_)
+
+    return (y_pred, test_labels)

birder/eval/methods/linear.py

@@ -0,0 +1,152 @@
+"""
+Linear probing on frozen embeddings
+
+Trains a single linear layer with cross-entropy loss for classification evaluation.
+"""
+
+import logging
+
+import numpy as np
+import numpy.typing as npt
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torch.utils.data import TensorDataset
+
+logger = logging.getLogger(__name__)
+
+
+def train_linear_probe(
+    train_features: npt.NDArray[np.float32],
+    train_labels: npt.NDArray[np.int_],
+    num_classes: int,
+    device: torch.device,
+    epochs: int,
+    batch_size: int,
+    lr: float = 1e-4,
+    step_size: int = 20,
+    seed: int = 0,
+) -> nn.Linear:
+    """
+    Train a linear probe on frozen embeddings
+
+    Parameters
+    ----------
+    train_features
+        Training features of shape (n_train, embedding_dim).
+    train_labels
+        Training labels of shape (n_train,), integer class indices.
+    num_classes
+        Number of output classes.
+    device
+        Device to train on.
+    epochs
+        Number of training epochs.
+    batch_size
+        Batch size for training.
+    lr
+        Learning rate.
+    step_size
+        Number of epochs between learning rate decay steps.
+    seed
+        Random seed for reproducibility.
+
+    Returns
+    -------
+    Trained linear layer.
+    """
+
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    input_dim = train_features.shape[1]
+    model = nn.Linear(input_dim, num_classes).to(device)
+
+    x_train = torch.from_numpy(train_features).float()
+    y_train = torch.from_numpy(train_labels).long()
+
+    dataset = TensorDataset(x_train, y_train)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.5)
+
+    model.train()
+    for epoch in range(epochs):
+        total_loss = 0.0
+        for batch_x, batch_y in loader:
+            batch_x = batch_x.to(device)
+            batch_y = batch_y.to(device)
+
+            optimizer.zero_grad()
+            logits = model(batch_x)
+            loss = criterion(logits, batch_y)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item() * batch_x.size(0)
+
+        scheduler.step()
+
+        if (epoch + 1) % 10 == 0:
+            avg_loss = total_loss / len(dataset)
+            logger.debug(f"Epoch {epoch + 1}/{epochs} - Loss: {avg_loss:.4f}")
+
+    return model
+
+
+def evaluate_linear_probe(
+    model: nn.Linear,
+    test_features: npt.NDArray[np.float32],
+    test_labels: npt.NDArray[np.int_],
+    batch_size: int = 128,
+    device: torch.device = torch.device("cpu"),
+) -> tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]:
+    """
+    Evaluate linear probe on test set
+
+    Parameters
+    ----------
+    model
+        Trained linear layer.
+    test_features
+        Test features of shape (n_test, embedding_dim).
+    test_labels
+        Test labels of shape (n_test,), integer class indices.
+    batch_size
+        Batch size for inference.
+    device
+        Device to run on.
+
+    Returns
+    -------
+    y_pred
+        Predicted labels for test samples.
+    y_true
+        True labels for test samples (same as test_labels).
+    """
+
+    x_test = torch.from_numpy(test_features).float()
+    y_test = torch.from_numpy(test_labels).long()
+
+    dataset = TensorDataset(x_test, y_test)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+    model.to(device)
+    model.eval()
+    all_preds: list[torch.Tensor] = []
+    all_labels: list[torch.Tensor] = []
+
+    with torch.inference_mode():
+        for batch_x, batch_y in loader:
+            batch_x = batch_x.to(device)
+            logits = model(batch_x)
+            preds = logits.argmax(dim=1)
+            all_preds.append(preds.cpu())
+            all_labels.append(batch_y)
+
+    y_pred = torch.concat(all_preds, dim=0).numpy().astype(np.int_)
+    y_true = torch.concat(all_labels, dim=0).numpy().astype(np.int_)
+
+    return (y_pred, y_true)