We recommend you to open this page on GitHub, so you can see the corresponding scripts and configs.
Pipelines: features extraction
Introduction to metric learning:
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:
extractor_training_pipeline including training + validation
extractor_validation_pipeline including validation only
extractor_prediction_pipeline including saving extracted features
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:
Accumulating all the embeddings in EmbeddingMetrics.
Calculating distances between queries and galleries.
[Optional] Applying some specific retrieval postprocessing techniques like re-ranking.
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 |
---|---|---|---|
|
|
Callable, see available. |
|
|
|
Callable, see available. |
|
|
|
A successor of IExtractor, see available. |
|
|
|
For losses with mining see this. For classification losses: no restrictions, but set |
|
|
|
The signature is required: |
|
|
|
A regular PyTorch optimizer. |
|
|
|
A regular PyTorch lr scheduler, structured in Lightning format. |
|
|
|
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 setn_instances
as big as your GPU allows for the givenn_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.