Package: pytorch-metric-learning
NewsJuly 25: v2.3.0AddedHistogramLossThank youdomenicoMuscill0.June 18: v2.2.0AddedManifoldLossandP2SGradLoss.Added asymmetricflag toSelfSupervisedLoss.See therelease notes.Thank youdomenicoMuscill0.DocumentationView the documentation hereView the installation instructions hereView the available losses, miners etc. hereGoogle Colab ExamplesSee theexamples folderfor notebooks you can download or run on Google Colab.PyTorch Metric Learning OverviewThis library contains 9 modules, each of which can be used independently within your existing codebase, or combined together for a complete train/test workflow.How loss functions workUsing losses and miners in your training loopLet’s initialize a plainTripletMarginLoss:frompytorch_metric_learningimportlossesloss_func=losses.TripletMarginLoss()To compute the loss in your training loop, pass in the embeddings computed by your model, and the corresponding labels. The embeddings should have size (N, embedding_size), and the labels should have size (N), where N is the batch size.# your training loopfori,(data,labels)inenumerate(dataloader):optimizer.zero_grad()embeddings=model(data)loss=loss_func(embeddings,labels)loss.backward()optimizer.step()The TripletMarginLoss computes all possible triplets within the batch, based on the labels you pass into it. Anchor-positive pairs are formed by embeddings that share the same label, and anchor-negative pairs are formed by embeddings that have different labels.Sometimes it can help to add a mining function:frompytorch_metric_learningimportminers,lossesminer=miners.MultiSimilarityMiner()loss_func=losses.TripletMarginLoss()# your training loopfori,(data,labels)inenumerate(dataloader):optimizer.zero_grad()embeddings=model(data)hard_pairs=miner(embeddings,labels)loss=loss_func(embeddings,labels,hard_pairs)loss.backward()optimizer.step()In the above code, the miner finds positive and negative pairs that it thinks are particularly difficult. Note that even though the TripletMarginLoss operates on triplets, it’s still possible to pass in pairs. This is because the library automatically converts pairs to triplets and triplets to pairs, when necessary.Customizing loss functionsLoss functions can be customized usingdistances,reducers, andregularizers. In the diagram below, a miner finds the indices of hard pairs within a batch. These are used to index into the distance matrix, computed by the distance object. For this diagram, the loss function is pair-based, so it computes a loss per pair. In addition, a regularizer has been supplied, so a regularization loss is computed for each embedding in the batch. The per-pair and per-element losses are passed to the reducer, which (in this diagram) only keeps losses with a high value. The averages are computed for the high-valued pair and element losses, and are then added together to obtain the final loss.Now here's an example of a customized TripletMarginLoss:frompytorch_metric_learning.distancesimportCosineSimilarityfrompytorch_metric_learning.reducersimportThresholdReducerfrompytorch_metric_learning.regularizersimportLpRegularizerfrompytorch_metric_learningimportlossesloss_func=losses.TripletMarginLoss(distance=CosineSimilarity(),reducer=ThresholdReducer(high=0.3),embedding_regularizer=LpRegularizer())This customized triplet loss has the following properties:The loss will be computed using cosine similarity instead of Euclidean distance.All triplet losses that are higher than 0.3 will be discarded.The embeddings will be L2 regularized.Using loss functions for unsupervised / self-supervised learningASelfSupervisedLosswrapper is provided for self-supervised learning:frompytorch_metric_learning.lossesimportSelfSupervisedLossloss_func=SelfSupervisedLoss(TripletMarginLoss())# your training for-loopfori,datainenumerate(dataloader):optimizer.zero_grad()embeddings=your_model(data)augmented=your_model(your_augmentation(data))loss=loss_func(embeddings,augmented)loss.backward()optimizer.step()If you're interested inMoCo-style self-supervision, take a look at theMoCo on CIFAR10notebook. It uses CrossBatchMemory to implement the momentum encoder queue, which means you can use any tuple loss, and any tuple miner to extract hard samples from the queue.Highlights of the rest of the libraryFor a convenient way to train your model, take a look at thetrainers.Want to test your model's accuracy on a dataset? Try thetesters.To compute the accuracy of an embedding space directly, useAccuracyCalculator.If you're short of time and want a complete train/test workflow, check out theexample Google Colab notebooks.To learn more about all of the above,see the documentation.InstallationRequired PyTorch versionpytorch-metric-learning >= v0.9.90requirestorch >= 1.6pytorch-metric-learning < v0.9.90doesn't have a version requirement, but was tested withtorch >= 1.2Other dependencies:numpy, scikit-learn, tqdm, torchvisionPippip install pytorch-metric-learningTo get the latest dev version:pip install pytorch-metric-learning --preTo install on Windows:pip install torch===1.6.0 torchvision===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install pytorch-metric-learningTo install with evaluation and logging capabilities(This will install the unofficial pypi version of faiss-gpu, plus record-keeper and tensorboard):pip install pytorch-metric-learning[with-hooks]To install with evaluation and logging capabilities (CPU)(This will install the unofficial pypi version of faiss-cpu, plus record-keeper and tensorboard):pip install pytorch-metric-learning[with-hooks-cpu]Condaconda install -c conda-forge pytorch-metric-learningTo use the testing module, you'll need faiss, which can be installed via conda as well. See theinstallation instructions for faiss.Benchmark resultsSeepowerful-benchmarkerto view benchmark results and to use the benchmarking tool.DevelopmentDevelopment is done on thedevbranch:git checkout devUnit tests can be run with the defaultunittestlibrary:python-munittestdiscoverYou can specify the test datatypes and test device as environment variables. For example, to test using float32 and float64 on the CPU:TEST_DTYPES=float32,float64TEST_DEVICE=cpupython-munittestdiscoverTo run a single test file instead of the entire test suite, specify the file name:python-munittesttests/losses/test_angular_loss.pyCode is formatted usingblackandisort:pipinstallblackisort
./format_code.shAcknowledgementsContributorsThanks to the contributors who made pull requests!ContributorHighlightsdomenicoMuscill0-ManifoldLoss-P2SGradLoss-HistogramLossmlopezantequera- Made thetesterswork on any combination of query and reference sets- MadeAccuracyCalculatorwork with arbitrary label comparisonscwkeam-SelfSupervisedLoss-VICRegLoss- Added mean reciprocal rank accuracy toAccuracyCalculator- BaseLossWrappermarijnl-BatchEasyHardMiner-TwoStreamMetricLoss-GlobalTwoStreamEmbeddingSpaceTester-Example using trainers.TwoStreamMetricLosschingisooinarSubCenterArcFaceLosselias-ramziHierarchicalSamplerfjsjSupConLossAlenUbuntuCircleLossinterestingzhuoPNPLosswconnellLearning a scRNAseq Metric EmbeddingAlexSchuyoptimizedutils.loss_and_miner_utils.get_random_triplet_indicesJohnGiorgiall_gatherinutils.distributedHummer12007utils.key_checkervltanhMadeInferenceModel.train_indexeraccept datasetsbtseytlinget_nearest_neighborsinInferenceModelmlw214Addedreturn_per_classtoAccuracyCalculatorlayumiInstanceLossNoTodyHelped addref_embandref_labelsto the distributed wrappers.ElisonShertonFixed an edge case in ArcFaceLoss.stompsjoImproved documentation for NTXentLossz1wthinline72tpanumfralikjoaqoJoOkumagkourosyutanakamura-tkyKinglittleQmartin0258michaeldeyzelHSinger04rheumbot66Facebook AIThank you toSer-Nam LimatFacebook AI, and my research advisor,Professor Serge Belongie. This project began during my internship at Facebook AI where I received valuable feedback from Ser-Nam, and his team of computer vision and machine learning engineers and research scientists. In particular, thanks toAshish ShahandAustin Reiterfor reviewing my code during its early stages of development.Open-source reposThis library contains code that has been adapted and modified from the following great open-source repos:https://github.com/bnu-wangxun/Deep_Metrichttps://github.com/chaoyuaw/incubator-mxnet/blob/master/example/gluon/embedding_learninghttps://github.com/facebookresearch/deepclusterhttps://github.com/geonm/proxy-anchor-losshttps://github.com/idstcv/SoftTriplehttps://github.com/kunhe/FastAP-metric-learninghttps://github.com/ronekko/deep_metric_learninghttps://github.com/tjddus9597/Proxy-Anchor-CVPR2020http://kaizhao.net/regularfacehttps://github.com/nii-yamagishilab/project-NN-Pytorch-scriptsLogoThanks toJeff Musgravefor designing the logo.Citing this libraryIf you'd like to cite pytorch-metric-learning in your paper, you can use this bibtex:@article{Musgrave2020PyTorchML,
  title={PyTorch Metric Learning},
  author={Kevin Musgrave and Serge J. Belongie and Ser-Nam Lim},
  journal={ArXiv},
  year={2020},
  volume={abs/2008.09164}}