Torchreid is a library for deep-learning person re-identification, written in PyTorch.
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.
Tech report: https://arxiv.org/abs/1910.10093.
You can find some research projects that are built on top of Torchreid here.
- [May 2020] Added the person attribute recognition code used in [Omni-Scale Feature Learning for Person Re-Identification (ICCV'19)](https://arxiv.org/abs/1905.00953).
- [May 2020]
1.2.1
: Added a simple API for feature extraction (torchreid/utils/feature_extractor.py
). See the documentation for the instruction. - [Apr 2020] Code for reproducing the experiments of deep mutual learning in the OSNet paper (Supp. B) has been released at
projects/DML
. - [Apr 2020] Upgraded to
1.2.0
. The engine class has been made more model-agnostic to improve extensibility. See Engine and ImageSoftmaxEngine for more details. Credit to Dassl.pytorch. - [Dec 2019] Our OSNet paper has been updated, with additional experiments (in section B of the supplementary) showing some useful techniques for improving OSNet's performance in practice.
- [Nov 2019]
ImageDataManager
can load training data from target datasets by settingload_train_targets=True
, and the train-loader can be accessed withtrain_loader_t = datamanager.train_loader_t
. This feature is useful for domain adaptation research.
Make sure 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
- Import
torchreid
import torchreid
- 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
)
- Build engine
engine = torchreid.engine.ImageSoftmaxEngine(
datamanager,
model,
optimizer=optimizer,
scheduler=scheduler,
label_smooth=True
)
- Run training and test
engine.run(
save_dir='log/resnet50',
max_epoch=60,
eval_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. The folder "configs/" contains some predefined configs which you can use as a starting point.
Below we provide an example to train and test OSNet (Zhou et al. ICCV'19). Assume PATH_TO_DATA
is the directory containing reid datasets. The environmental variable CUDA_VISIBLE_DEVICES
is omitted, which you need to specify if you have a pool of gpus and want to use a specific set of them.
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
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 \
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 \
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
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.
- OSNet-IBN1-Lite (test-only code with lite docker container)
- Deep Learning for Person Re-identification: A Survey and Outlook
If you find this code useful to your research, please cite the following papers.
@article{torchreid,
title={Torchreid: A Library for Deep Learning Person Re-Identification in Pytorch},
author={Zhou, Kaiyang and Xiang, Tao},
journal={arXiv preprint arXiv:1910.10093},
year={2019}
}
@inproceedings{zhou2019osnet,
title={Omni-Scale Feature Learning for Person Re-Identification},
author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
booktitle={ICCV},
year={2019}
}
@article{zhou2019learning,
title={Learning Generalisable Omni-Scale Representations for Person Re-Identification},
author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
journal={arXiv preprint arXiv:1910.06827},
year={2019}
}