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
Documentation: https://kaiyangzhou.github.io/deep-person-reid/.
Code: https://github.com/KaiyangZhou/deep-person-reid.
Model zoo: https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.
The code works with both python2 and python3.
- Install PyTorch and torchvision following the official instructions.
- Clone
deep-person-reid
to your preferred directory
$ git clone https://github.com/KaiyangZhou/deep-person-reid.git
cd
todeep-person-reid
and install dependencies
$ cd deep-person-reid/
$ pip install -r requirements.txt
- 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
We also provide an environment.yml file for easy setup with conda.
- Clone
deep-person-reid
to your preferred directory
$ git clone https://github.com/KaiyangZhou/deep-person-reid.git
cd
todeep-person-reid
and create an environment (namedtorchreid
)
$ cd deep-person-reid/
$ conda env create -f environment.yml
In doing so, the dependencies will be automatically installed.
- 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
- 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
- Import
torchreid
import torchreid
- Load data manager
datamanager = torchreid.data.ImageDataManager(
root='reid-data',
sources='market1501',
height=256,
width=128,
batch_size=32,
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,
print_freq=10,
test_only=False
)
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 \
--transforms random_flip random_crop \
--save-dir log/resnet50-market1501-softmax \
--gpu-devices 0
Please refer to default_parser.py
and main.py
for more details.
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}
}