ml-lab / magnetloss-pytorch

PyTorch implementation of a deep metric learning technique called "Magnet Loss" from Facebook AI Research (FAIR).

Magnet-Loss-PyTorch

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.

Table of Contents

• Installation
• Anaconda
• Docker
• Results
• Other Implementations
• Citing MagnetLoss-PyTorch

Installation

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

Anaconda

Anaconda: Installation

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

Anaconda: Train

Train ConvNet with Magnet Loss on the local machine using MNIST dataset:

python magnet_loss_test.py --lr 1e-4 --batch-size 64 --mnist --dml

Docker GPU Training

Prerequisites:

1. Docker installed on your machine. If you don't have Docker installed already, then go here to Docker Setup
2. Install nvidia-docker 2.0 from Nvidia Docker 2.0
3. Register nvidia runtime with the Docker engine using Nvidia Container Runtime

Docker: Build Image

docker build -t magnetloss .

Docker: Train

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 --dml

or

./run_gpu_docker.sh magnetloss

Results

MNIST

Iterations	Learned Embedding Space
0
2000
4000
6000
8000
10000
12000
14000

Other Implementations

• Magnet Loss in Tensorflow

Citing MagnetLoss-PyTorch

If you use MagnetLoss-PyTorch in a scientific publication, I would appreciate references to the source code.

Biblatex entry:

@misc{vithu2018,
  author = {Thangarasa, Vithursan},
  title = {MagnetLoss-PyTorch},
  year = {2018},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/vithursant/MagnetLoss-PyTorch}}
}