huizhang0110 / deep-person-reid

Deep learning person re-identification in PyTorch.

Torchreid

Torchreid is a library built on PyTorch for deep-learning person re-identification.

It features:

• multi-GPU training
• support both image- and video-reid
• end-to-end training and evaluation
• incredibly easy preparation of reid datasets
• multi-dataset training
• cross-dataset evaluation
• standard protocol used by most research papers
• highly extensible (easy to add models, datasets, training methods, etc.)
• implementations of state-of-the-art deep reid models
• access to pretrained reid models
• advanced training techniques
• visualization of ranking results

Documentation: https://kaiyangzhou.github.io/deep-person-reid/.

News

• 09-05-2019: The person re-ranking code has been added to this repo.
• 06-05-2019: We released a tech report on arxiv. Code and models will be released.
• 24-03-2019: Torchreid documentation is out!

Installation

The code works with both python2 and python3.

Option 1

1. Install PyTorch and torchvision following the official instructions.
2. Clone deep-person-reid to your preferred directory

$ git clone https://github.com/KaiyangZhou/deep-person-reid.git

3. cd to deep-person-reid and install dependencies

$ cd deep-person-reid/
$ pip install -r requirements.txt

4. Install torchreid

$ python setup.py install # or python3
$ # If you wanna modify the source code without
$ # the need to rebuild it, you can do
$ # python setup.py develop

Option 2 (with conda)

We also provide an environment.yml file for easy setup with conda.

1. Clone deep-person-reid to your preferred directory

$ git clone https://github.com/KaiyangZhou/deep-person-reid.git

2. cd to deep-person-reid and create an environment (named torchreid)

$ cd deep-person-reid/
$ conda env create -f environment.yml

In doing so, the dependencies will be automatically installed.

3. Install PyTorch and torchvision (select the proper cuda version to suit your machine)

$ conda activate torchreid
$ conda install pytorch torchvision cudatoolkit=9.0 -c pytorch

4. Install torchreid

$ python setup.py install
$ # If you wanna modify the source code without
$ # the need to rebuild it, you can do
$ # python setup.py develop

Get started: 30 seconds to Torchreid

1. Import torchreid

import torchreid

2. Load data manager

datamanager = torchreid.data.ImageDataManager(
    root='reid-data',
    sources='market1501',
    height=256,
    width=128,
    batch_size=32,
    market1501_500k=False
)

3 Build model, optimizer and lr_scheduler

model = torchreid.models.build_model(
    name='resnet50',
    num_classes=datamanager.num_train_pids,
    loss='softmax',
    pretrained=True
)

model = model.cuda()

optimizer = torchreid.optim.build_optimizer(
    model,
    optim='adam',
    lr=0.0003
)

scheduler = torchreid.optim.build_lr_scheduler(
    optimizer,
    lr_scheduler='single_step',
    stepsize=20
)

4. Build engine

engine = torchreid.engine.ImageSoftmaxEngine(
    datamanager,
    model,
    optimizer=optimizer,
    scheduler=scheduler,
    label_smooth=True
)

5. Run training and test

engine.run(
    save_dir='log/resnet50',
    max_epoch=60,
    eval_freq=10,
    print_freq=10,
    test_only=False
)

A unified interface

In "deep-person-reid/scripts/", we provide a unified interface including a default parser file default_parser.py and the main script main.py. For example, to train an image reid model on Market1501 using softmax, you can do

python main.py \
--root path/to/reid-data \
--app image \
--loss softmax \
--label-smooth \
-s market1501 \
-a resnet50 \
--optim adam \
--lr 0.0003 \
--max-epoch 60 \
--stepsize 20 40 \
--batch-size 32 \
--save-dir log/resnet50-market-softmax \
--gpu-devices 0

Please refer to default_parser.py and main.py for more details.

Datasets

Image-reid datasets

• Market1501
• CUHK03
• DukeMTMC-reID
• MSMT17
• VIPeR
• GRID
• CUHK01
• SenseReID
• QMUL-iLIDS
• PRID

Video-reid datasets

• MARS
• iLIDS-VID
• PRID2011
• DukeMTMC-VideoReID

Models

ImageNet classification models

• ResNet
• ResNeXt
• SENet
• DenseNet
• Inception-ResNet-V2
• Inception-V4
• Xception

Lightweight models

• NASNet
• MobileNetV2
• ShuffleNet
• SqueezeNet

ReID-specific models

• MuDeep
• ResNet-mid
• HACNN
• PCB
• MLFN

Losses

• Softmax (cross entropy loss with label smoothing)
• Triplet (hard example mining triplet loss)

Citation

If you find this code useful to your research, please cite the following publication.

@article{zhou2019osnet,
title={Omni-Scale Feature Learning for Person Re-Identification},
author={Kaiyang Zhou and Yongxin Yang and Andrea Cavallaro and Tao Xiang},
journal={arXiv preprint arXiv:1905.00953},
year={2019}
}