Skip to content

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

License

Notifications You must be signed in to change notification settings

fengshi-cherish/MagnetLoss-PyTorch

 
 

Repository files navigation

MagnetLoss-PyTorch

PyTorch implementation of the Magnet Loss for Deep Metric Learning, based on the following paper:

Table of Contents

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 --magnet-loss

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 --magnet-loss

or

./run_gpu_docker.sh magnetloss

Results

MNIST

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

Citing MagnetLoss-PyTorch

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

About

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

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 99.3%
  • Shell 0.7%