PyTorch implementation of the Magnet Loss for Deep Metric Learning, based on the following paper:
- Metric Learning with Adaptive Density Discrimination by Oren Rippel, Piotr Dollar, Manohar Paluri, Lubomir Bourdev from Facebook AI Research that was accepted into ICLR 2016.
The program requires the following dependencies (easy to install using pip, Ananconda or Docker):
- python (tested on 2.7 and 3.6)
- pytorch (tested with v0.3 and v0.3.1 with CUDA 8.0/9.0)
- numpy
- matplotlib
- seaborn
- pandas
- tqdm
- pillow
- sklearn
- scipy
- visdom
To install MagnetLoss in an Anaconda environment:
(Python 2.7) conda env create -f pytorch-2p7-cuda80.yml
(Python 3.6) conda env create -f pytorch-3p6-cuda80.yml
To activate Anaconda environment:
(Python 2.7) source activate magnet-loss-py27-env
(Python 3.6) source activate magnet-loss-py36-env
Train ConvNet with Magnet Loss on the local machine using MNIST dataset:
python magnet_loss_test.py --lr 1e-4 --batch-size 64 --mnist --magnet-loss
- Docker installed on your machine. If you don't have Docker installed already, then go here to Docker Setup
- Install
nvidia-docker 2.0
from Nvidia Docker 2.0 - Register
nvidia
runtime with the Docker engine using Nvidia Container Runtime
docker build -t magnetloss .
Deploy and train on Docker container:
docker run --rm -it --runtime=nvidia magnetloss python magnet_loss_test.py --lr 1e-4 --mnist --batch-size 64 --magnet-loss
or
./run_gpu_docker.sh magnetloss
Iterations | Learned Embedding Space |
---|---|
0 | |
2000 | |
4000 | |
6000 | |
8000 | |
10000 | |
12000 | |
14000 |
If you use MagnetLoss-PyTorch in a scientific publication, I would appreciate references to the source code.
Biblatex entry:
@misc{MagnetLossPyTorch,
author = {Thangarasa, Vithursan},
title = {MagnetLoss-PyTorch},
year = {2018},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/vithursant/MagnetLoss-PyTorch}}
}