DOCSEARCHbyAlgolia
Submit Docs

Open Metric Learning (OML)

repository·main·Indexed 21 days ago

https://github.com/oml-team/open-metric-learning

A PyTorch-based framework for training and validating models that produce high-quality embeddings for metric learning tasks. OML provides modular pipelines for feature extraction, validation, and inference, supporting contrastive losses (e.g., TripletLoss) and classification losses (e.g., ArcFace). It includes tools for preparing benchmark datasets such as CARS 196, CUB 200 2011, InShop, and SOP, as well as a postprocessing pipeline for Siamese Transformer for Image Retrieval (STIR).

Tokens
42.1K
Snippets
108
Records
160
Agent score
76%

What's inside Open Metric Learning

  1. Overview of Open Metric Learning (OML)

    main
    OML (Open Metric Learning) is a PyTorch-based framework designed for training and validating models that produce high-quality embeddings. It provides tools for metric learning tasks, enabling developers to build models capable of learning meaningful representations in a shared embedding space.

  2. What is Open Metric Learning (OML)?

    main

    Open Metric Learning (OML) is a Python-based framework designed for training and validating deep learning models that produce high-quality embeddings for metric learning tasks (also known as extreme classification).

    Unlike vanilla classifiers, OML is optimized for scenarios where you have many entities (thousands of IDs) but few samples per entity, and where the goal is often retrieval or searching (e.g., Face Recognition, Re-Identification, or Product Search).

    Key differentiators from other libraries like PyTorch Metric Learning (PML) include:

    • Pipelines: Config-based training workflows that handle data preparation and model training end-to-end.
    • Zoo: A collection of pretrained models accessible similarly to torchvision.
    • PyTorch Lightning Integration: Built-in support for high-performance training loops and Distributed Data Parallel (DDP).
    • Recipe-Oriented: Focuses on practical, real-world benchmarks and hyperparameter combinations rather than just providing a collection of low-level tools.

  4. What is STIR (Siamese Transformer for Image Retrieval)?

    main
    STIR is a pairwise postprocessing (re-ranking) approach designed for image retrieval. Unlike traditional re-ranking transformers that rely on global or local feature extraction, STIR uses an attention mechanism to directly compare a query image and a retrieved candidate at the pixel level in a single forward pass. This method is designed to re-rank the top outputs of a baseline model (typically trained with triplet loss and hard negative mining) to achieve state-of-the-art performance on datasets like Stanford Online Products and DeepFashion In-shop.

  5. What is logged during OML Pipeline execution

    main

    When using integrated loggers in Pipelines, the following information is automatically captured:

    • Metrics: Performance metrics defined in your config via metric_args (e.g., CMC@1, Precision@5, MAP@5).
      • Tip: Set metrics_args.return_only_overall_category: False to log metrics independently for each category in multi-category datasets.
    • Loss Values: Averaged over batches and epochs. Some built-in OML losses (like TripletLossWithMargin) log additional statistics such as positive/negative distances and the fraction of active triplets.
    • Visualizations: Error analysis visualizations (e.g., showing queries, closest irrelevant items, and ground truths).
    • Reproducibility Artifacts: Source code, configuration files, dataframes, and tags.

  6. Understand Feature Extraction Pipelines

    main

    Feature extraction pipelines in OML are designed to train, validate, and perform inference on models that represent images as feature vectors. There are three primary pipeline types:

    1. extractor_training_pipeline: Handles both training and validation.
    2. extractor_validation_pipeline: Performs validation only.
    3. extractor_prediction_pipeline: Runs inference on a trained model and saves the extracted features to disk.

    These pipelines can be used via configuration files or as plain Python implementations.

  7. Understand the core concepts and terminology of OML

    main

    To use OML effectively, it is important to understand its specific naming conventions and the metric learning problem it solves:

    Metric Learning Problem

    Metric learning (or extreme classification) involves training models where you have thousands of unique IDs (entities) but only a few samples per ID. The goal is to produce embeddings that allow for efficient searching or matching, often involving unseen entities at test time.

    Glossary

    • embedding: The model's output (also called features vector or descriptor).
    • query: A sample used as a request in a retrieval procedure.
    • gallery set: The set of entities to search through (also called reference or index).
    • Sampler: A DataLoader argument used to form batches.
    • Miner: An object that forms pairs or triplets after the Sampler has formed a batch. Miners may use a memory bank to find combinations outside the current batch.
    • Samples/Labels/Instances: In a dataset, labels are the unique IDs (e.g., fashion item IDs), and instances or samples are the individual photos/data points for those labels.
    • categories: Groups of labels (e.g., "jackets").
    • training epoch: In OML, an epoch typically refers to observing all available labels rather than all available samples, because batch samplers for combination-based losses are often defined by the number of labels.

  8. Access pretrained models from the OML Zoo

    main
    OML provides a Zoo of pretrained models. These can be accessed in your code in a manner similar to torchvision. This is useful for both direct use and as a starting point for further training (e.g., initializing a model with a self-supervised checkpoint instead of standard ImageNet weights).

  9. How Pipelines are structured

    main

    Every OML Pipeline is composed of three fundamental building blocks:

    1. Config file (.yaml): Describes the entire experiment run, including model parameters, device settings, and component selections.
    2. Registry of classes & functions: A mapping system that connects configuration keys to Python constructors or functions.
    3. Entrypoint function & script: The logic that orchestrates the pipeline execution, typically invoked via a script using Hydra for configuration management.

    Because Pipelines use Hydra, you can override configuration values directly from the command line (e.g., python script.py model.args.weights=null).

    python validate.py model.args.weights=null

  10. Boost retrieval accuracy using a pairwise model as a re-ranker

    main
    You can improve the accuracy of vector search results by using a Siamese (pairwise) model as a re-ranker. This process involves taking the top-$N$ outputs from an initial retrieval model and performing inference on pairs consisting of the query and each retrieved output (query_i, output_j) where j=1..top_n. The pairwise model then re-ranks these outputs based on their similarity scores.

  11. Use bounding boxes and audio offsets in datasets

    main

    OML supports additional metadata in the CSV format for specialized tasks:

    • Bounding Boxes (Images): Use x_1, x_2, y_1, and y_2 columns. The format is left, right, top, bot. Ensure that x_1 < x_2 and y_1 < y_2. If only some images have boxes, leave the other rows empty.
    • Audio Offsets (Audios): Use the start_time column (float) to specify the offset from which the audio should start being read.
    • Re-identification (Re-id): Use the sequence column (string or integer) to group sequences of photos.

  12. Configure Training with Contrastive vs Classification Losses

    main

    When setting up a training pipeline, the choice of loss function dictates your sampling and mining requirements:

    Contrastive Losses (e.g., TripletLoss)

    • Requirement: Requires a Miner (to produce triplets using strategies like hard mining) and a Batches Sampler (to ensure batches contain enough different labels to form triplets).
    • Mechanism: Mining occurs inside the forward(features, labels) pass.

    Classification Losses (e.g., ArcFace)

    • Requirement: No mining step is required by design. Batch sampling is optional.
    • Mechanism: The feature vector is considered to be the output of the layer immediately preceding the classification head. Mining is not used, and the forward(features, labels) signature is still maintained.