基于pytorch的深度学习人员重识别
Torchreid是一个库。基于pytorch的深度学习人员重识别。
特点:
支持多GPU训练 支持图像的人员重识别与视频的人员重识别 端到端的训练与评估 简单的reid数据集的预准备 多种数据集训练 交叉数据集评估 支持标准协议 高扩展性 (容易追加模型,数据集, 训练方法等.) 实现了state-of-the-art deep reid models 预训练reid模型 高级训练技巧 可视化工具 (tensorboard, ranks, etc.)
安装:
下载代码后,
cd deep-person-reid/ conda create --name torchreid python=3.7 conda activate torchreid
pip install -r requirements.txt
conda install pytorch torchvision cudatoolkit=9.0 -c pytorch
python setup.py develop
使用:
加载torchreid,数据,然后构建模型。
import torchreid
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'])
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)
engine = torchreid.engine.ImageSoftmaxEngine( datamanager, model, optimizer=optimizer, scheduler=scheduler, label_smooth=True)
训练及评估
engine.run( save_dir='log/resnet50', max_epoch=60, eval_freq=10, print_freq=10, test_only=False)
使用的图像reid数据集有Market1501,CUHK03,DukeMTMC-reID MSMT17,VIPeR,GRID,CUHK01,SenseReID,QMUL-iLIDS,PRID。
使用的视频reid数据集有MARS,iLIDS-VID,PRID2011,DukeMTMC-VideoReID。