Embeddings Models

[kozlov721]

Overview

Re-opening of #141

• Added GhostFaceNetV2 network for embeddings learning together with corresponding losses, metrics and visualizer
  • ClosestIsPositiveAccuracy and MedianDistance metrics
  • Most losses from pytorch-metric-learning
  • EmbeddingVisualizer
• Integrated the "metadata" labels introduced in #211
• Groundwork for type checking of the metadata labels and better automatic inference of correct tasks

Example:

from luxonis_train.enums import Metadata

class Node(BaseNode):
    tasks = [Metadata("id")]

class Loss(BaseLoss):
    supported_tasks = [Metadata("id")]

The specific name of the metadata field can be overridden in the config file using metadata_task_override field.

- name: GhostFaceNetHead
  alias: color-embeddings
  metadata_task_override: color

This will cause the model to look for "task_name/metadata/color" labels instead of "task_name/metadata/id"

Open questions

Predefined model

There are a lot of options to specify for the embedding models, we support:

• 21 losses
• 5 distance measures
• 10 miners
• 9 reducers
• 5 regularizers

It would be good to do a large benchmarking on some real dataset to determine good combinations that we can use in the predefined models

Visualizer

Visualizing the embedding results is challenging. The current implementation uses PCA to reduce the dimensionality to 2d and then plot the points using seaborn.scatterplot and seaborn.kdeplot so we can see whether the points representing different embedding classes are getting grouped together.

This works somewhat good for tasks with just a few embedding classes, but is not great for anything more complex

• The plots get very convoluted and not really readable
• Visualizers are called per batch, but with many embedding classes the batch will usually contain only one example per class
  • We can use a little hack to accumulate the batches (the commented code in the PR)
  • We could generalize the visualizers so they work like metrics with reset, update, and compute methods which would allow natural batch accumulation
    • This would also allow visualizing closest matches using database and query images
    • Could be eventually useful for other visualizers as well

Comparison of simple embeddings of vehicle images based on color (red, green, blue)

Untrained model on top, trained on the bottom

Dataset with many classes (VeRI dataset)

Untrained model on the right, trained on the left

[klemen1999]

Generally LGTM

On the open questions:

Predefined Model

• I think currently predefined model can use the same setup as it is definened in the embeddings_model.yaml (same loss, miner, distance, reducer, etc.). This was used before already for a successful training and works. But agree, we should also test out others and give some pointers in the docs which combination is better for specific cases - this can be addressed as a separate PR IMO

Visualizer

I like the second idea of generalizing visualizers to same interface as we have for metrics. I agree that longterm this would be the best solution. I wouldn't make this a blocking requirement but we can note internally as a nice to have improvement.

And for the visualization of embedding models for datasets with a lot of classes I think to some degree this is still unsolved problem. There are multiple methods for reducing dimensionality (PCA, t-SNE, UMAP, etc.) and we likely can't catch all as the best one is usually decided case by case. But we could provide several options to choose from. But also nice-to-have, for initial implementation current one is nice.