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 tools (tensorboard, ranks, etc.)

Code: https://github.com/KaiyangZhou/deep-person-reid.

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

How-to instructions: https://kaiyangzhou.github.io/deep-person-reid/user_guide.

Model zoo: https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.

Make sure your conda is installed.

# cd to your preferred directory and clone this repo
git clone https://github.com/KaiyangZhou/deep-person-reid.git

# create environment
cd deep-person-reid/
conda create --name torchreid python=3.7
conda activate torchreid

# install dependencies
# make sure `which python` and `which pip` point to the correct path
pip install -r requirements.txt

# install torch and torchvision (select the proper cuda version to suit your machine)
conda install pytorch torchvision cudatoolkit=9.0 -c pytorch

# install torchreid (don't need to re-build it if you modify the source code)
python setup.py develop

1. Import torchreid

import torchreid

2. Load data manager

datamanager = torchreid.data.ImageDataManager(
    root='reid-data',
    sources='market1501',
    targets='market1501',
    height=256,
    width=128,
    batch_size_train=32,
    batch_size_test=100,
    transforms=['random_flip', 'random_crop']
)

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
)

In "deep-person-reid/scripts/", we provide a unified interface to train and test a model. See "scripts/main.py" and "scripts/default_config.py" for more details. "configs/" contains some predefined configs which you can use as a starting point.

Below we provide examples to train and test OSNet (Zhou et al. ICCV'19). Assume PATH_TO_DATA is the directory containing reid datasets.

To train OSNet on Market1501, do

python scripts/main.py \
--config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \
--transforms random_flip random_erase \
--root $PATH_TO_DATA \
--gpu-devices 0

The config file sets Market1501 as the default dataset. If you wanna use DukeMTMC-reID, do

python scripts/main.py \
--config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \
-s dukemtmcreid \
-t dukemtmcreid \
--transforms random_flip random_erase \
--root $PATH_TO_DATA \
--gpu-devices 0 \
data.save_dir log/osnet_x1_0_dukemtmcreid_softmax_cosinelr

The code will automatically (download and) load the ImageNet pretrained weights. After the training is done, the model will be saved as "log/osnet_x1_0_market1501_softmax_cosinelr/model.pth.tar-250". Under the same folder, you can find the tensorboard file. To visualize the learning curves using tensorboard, you can run tensorboard --logdir=log/osnet_x1_0_market1501_softmax_cosinelr in the terminal and visit http://localhost:6006/ in your web browser.

Evaluation is automatically performed at the end of training. To run the test again using the trained model, do

python scripts/main.py \
--config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \
--root $PATH_TO_DATA \
--gpu-devices 0 \
model.load_weights log/osnet_x1_0_market1501_softmax_cosinelr/model.pth.tar-250 \
test.evaluate True

Suppose you wanna train OSNet on DukeMTMC-reID and test its performance on Market1501, you can do

python scripts/main.py \
--config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad.yaml \
-s dukemtmcreid \
-t market1501 \
--transforms random_flip color_jitter \
--root $PATH_TO_DATA \
--gpu-devices 0

Here we only test the cross-domain performance. However, if you also want to test the performance on the source dataset, i.e. DukeMTMC-reID, you can set -t dukemtmcreid market1501, which will evaluate the model on the two datasets separately.

Different from the same-domain setting, here we replace random_erase with color_jitter. This can improve the generalization performance on the unseen target dataset.

Pretrained models are available in the Model Zoo.

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

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

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

• NASNet
• MobileNetV2
• ShuffleNet
• ShuffleNetV2
• SqueezeNet

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

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

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={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
  journal={arXiv preprint arXiv:1905.00953},
  year={2019}
}