A flexible and extensible metric learning library, written in PyTorch.
Project description
pytorch_metric_learning
Coming soon in a separate repo: a highly configurable and easy-to-use benchmarking tool.
Loss functions implemented:
- angular
- contrastive
- lifted structure
- margin
- multi similarity
- n pairs
- nca
- proxy nca
- triplet margin
- more to be added
Mining functions implemented:
- distance weighted sampling
- hard aware cascaded mining
- maximum loss miner
- multi similarity miner
- pair margin miner
- more to be added
Training methods implemented:
- metric loss only
- training with classifier
- cascaded embeddings
- deep adversarial metric learning
- more to be added
Installation:
pip install pytorch_metric_learning
Overview
Use a loss function by itself
from pytorch_metric_learning import losses
loss_func = losses.TripletMarginLoss(normalize_embeddings=False, margin=0.1)
loss = loss_func(embeddings, labels)
Or combine miners and loss functions, regardless of whether they mine or compute loss using pairs or triplets. Pairs are converted to triplets when necessary, and vice versa.
from pytorch_metric_learning import miners, losses
miner = miners.MultiSimilarityMiner(epsilon=0.1)
loss_func = losses.TripletMarginLoss(normalize_embeddings=False, margin=0.1)
hard_pairs = miner(embeddings, labels)
loss = loss_func(embeddings, labels, hard_pairs)
Train using more advanced approaches, like deep adversarial metric learning. For example:
from pytorch_metric_learning import trainers
# Set up your models, optimizers, loss functions etc.
models = {"trunk": your_trunk_model,
"embedder": your_embedder_model,
"G_neg_model": your_negative_generator}
optimizers = {"trunk_optimizer": your_trunk_optimizer,
"embedder_optimizer": your_embedder_optimizer,
"G_neg_model_optimizer": your_negative_generator_optimizer}
loss_funcs = {"metric_loss": losses.AngularNPairs(alpha=35),
"synth_loss": losses.Angular(alpha=35),
"G_neg_adv": losses.Angular(alpha=35)}
mining_funcs = {}
loss_weights = {"metric_loss": 1,
"classifier_loss": 0,
"synth_loss": 0.1,
"G_neg_adv": 0.1,
"G_neg_hard": 0.1,
"G_neg_reg": 0.1}
# Create trainer object
trainer = trainers.DeepAdversarialMetricLearning(
models=models,
optimizers=optimizers,
batch_size=120,
loss_funcs=loss_funcs,
mining_funcs=mining_funcs,
num_epochs=50,
iterations_per_epoch=100,
dataset=your_dataset,
loss_weights=loss_weights
)
# Train!
trainer.train()
The package also comes with RecordKeeper, which makes it very easy to log and save data during training. It automatically looks for special attributes within objects to log on Tensorboard, as well as to save in CSV and pickle format.
from torch.utils.tensorboard import SummaryWriter
from pytorch_metric_learning.utils import record_keeper as record_keeper_package
pickler_and_csver = record_keeper_package.PicklerAndCSVer(your_folder_for_logs)
tensorboard_writer = SummaryWriter(log_dir=your_tensorboard_folder)
record_keeper = record_keeper_package.RecordKeeper(tensorboard_writer, pickler_and_csver)
# Then during training:
recorder.update_records(your_dict_of_objects, current_iteration)
# If you are using one of the provided trainers, then just pass in the record keeper, and the update step will be taken care of.
trainer = trainers.MetricLossOnly(
<your other args>
record_keeper = record_keeper
...
)
# Now it will update the record_keeper at every iteration
trainer.train()
The nice thing about RecordKeeper is that it makes it very easy to add loggable information when you write a new loss function or miner. Just create a list named "record_these" that contains the names of the attributes you want to record.
class YourNewLossFunction(BaseMetricLossFunction):
def __init__(self, **kwargs):
self.avg_embedding_norm = 0
self.some_other_useful_stat = 0
self.record_these = ["avg_embedding_norm", "some_other_useful_stat"]
super().__init__(**kwargs)
def compute_loss(self, embeddings, labels, indices_tuple):
self.avg_embedding_norm = torch.mean(torch.norm(embeddings, p=2, dim=1))
self.some_other_useful_stat = some_cool_function(embeddings)
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for pytorch_metric_learning-0.9.14.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 387bf61a3832fad1749eeb368090ce4595ceb6ad217e1c761ee496527398dcc6 |
|
MD5 | 6fbb0cc5fa8a2dcc958958119f9cf910 |
|
BLAKE2b-256 | c4089c2d84fcef39c5685fd89197d49da26bcb42303b173386a5f0c2934257c6 |
Hashes for pytorch_metric_learning-0.9.14-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | fca62aeeddf0f9085fc33c604728e618341dec0be3e7f77aa8b4cc6945095900 |
|
MD5 | 415cf5161936a0a0304c867721169004 |
|
BLAKE2b-256 | d3eb069f3868b23982710d715782d2ac4953787fc8758fb87d4c1ff36305eac9 |