We recommend you to open this page on GitHub, so you can see the corresponding scripts and configs.

Pipelines: features extraction

English | Russian | Chinese

These particular pipelines allow you to train, validate and inference models that represent images as feature vectors. In this section we explain how the following pipelines work under the hood:

Pipelines also have the corresponding analogues in plain Python.

Training

It is expected that the dataset will be in the desired format. You can see a tiny figures dataset as an example.

To get used to terminology you can check the Glossary (naming convention).

This pipeline support two types of losses:

  • Contrastive ones, like TripletLoss. They require special Miner and Batches Sampler. Miner produces triplets exploiting different strategies like hard mining, in its turn Sampler guarantees that batch contains enough different labels to form at least one triplet so miner can do its job.

  • Classification ones, like ArcFace. They have no mining step by design and batch sampling strategy is optional for them. For these losses we consider the output of the layer before the classification head (which is a part of criterion in our implementation) as a feature vector.

Note! Despite the different nature of the losses above, they share the same forward signature: forward(features, labels). That is why mining is happening inside the forward pass, see TripletLossWithMiner.

Validation

Validation part consists of the following steps:

  1. Accumulating all the embeddings in EmbeddingMetrics.

  2. Calculating distances between queries and galleries.

  3. [Optional] Applying some specific retrieval postprocessing techniques like re-ranking.

  4. Calculating retrieval metrics like CMC@k, Precision@k, MeanAveragePrecision@k or others.

Prediction / Inference

Prediction pipeline runs inference of a trained model and saves extracted features to the disk. Note, to speed up inference you can easily turn on multi GPU setup in the corresponding config file.

Customization

Pipelines are built around blocks like model, criterion, optimizer and so on. Some of them can be replaced by existing entities from OML or by your custom implementations, see the customisation instruction.

In feature extraction pipelines you can customize:

Block in config

Registry*

Example configs

Requirements on custom implementation

transforms_train

TRANSFORMS_REGISTRY

configs

Callable, see available.

transforms_val

TRANSFORMS_REGISTRY

configs

Callable, see available.

extractor

EXTRACTORS_REGISTRY

configs

A successor of IExtractor, see available.

sampler

SAMPLERS_REGISTRY

configs

For losses with mining see this. For classification losses: no restrictions, but set null for RandomSampler.

criterion

LOSSES_REGISTRY

configs

The signature is required: forward(features, labels). For contrastive losses: mining is implemented inside the forward pass. For classification losses: a classification head is a part of criterion. See available.

optimizer

OPTIMIZERS_REGISTRY

configs

A regular PyTorch optimizer.

scheduling

SCHEDULERS_REGISTRY

configs

A regular PyTorch lr scheduler, structured in Lightning format.

logger

LOGGERS_REGISTRY

configs

Child of IPipelineLogger

*Use: from oml.registry import X_REGISTRY.

Tips

We left plenty of comments in the training config for the CARS dataset, so you can start checking it out.

  • If you don’t know what parameters to pick for BalanceSampler, simply set n_labels equal to the median size of your classes, and set n_instances as big as your GPU allows for the given n_labels.

  • Tips for TripletLossWithMiner:

    • The margin value of 0.2 may be a good choice if your extractor produces normalised features.

    • Triplet loss may struggle with a mode collapse: the situation when your loss goes down, but then fluctuates on a plateau on the level of margin value, which means that positive and negative distances both equal to zero. In this case, you can try to use the soft version of triplet loss instead (just set margin: null). You can also switch between mining strategies (hard / all).

    • Don’t use margin: null if you normalise features since it breaks gradients flow (it can be proved mathematically).

  • Check out Logging & Visualization to learn more about built-in possibilities such as tracking metrics, losses, geometric statistics over embeddings, visual inspection of a model’s predictions and so on.