├── .idea ├── .gitignore ├── TransReID-main.iml ├── deployment.xml ├── inspectionProfiles │ ├── Project_Default.xml │ └── profiles_settings.xml ├── misc.xml ├── modules.xml └── remote-mappings.xml ├── LICENSE ├── README.md ├── config ├── __init__.py └── defaults.py ├── configs ├── Cuhk03_labeled │ └── vit_transreid_stride.yml ├── DukeMTMC │ ├── deit_transreid_stride.yml │ ├── vit_base.yml │ ├── vit_jpm.yml │ ├── vit_sie.yml │ ├── vit_transreid.yml │ ├── vit_transreid_384.yml │ ├── vit_transreid_stride.yml │ └── vit_transreid_stride_384.yml ├── MSMT17 │ ├── deit_small.yml │ ├── deit_transreid_stride.yml │ ├── vit_base.yml │ ├── vit_jpm.yml │ ├── vit_sie.yml │ ├── vit_small.yml │ ├── vit_transreid.yml │ ├── vit_transreid_384.yml │ ├── vit_transreid_stride.yml │ └── vit_transreid_stride_384.yml ├── Market │ ├── deit_transreid_stride.yml │ ├── vit_base.yml │ ├── vit_jpm.yml │ ├── vit_sie.yml │ ├── vit_transreid.yml │ ├── vit_transreid_384.yml │ ├── vit_transreid_stride.yml │ └── vit_transreid_stride_384.yml ├── OCC_Duke │ ├── deit_transreid_stride.yml │ ├── vit_base.yml │ ├── vit_jpm.yml │ ├── vit_sie.yml │ ├── vit_transreid.yml │ └── vit_transreid_stride.yml ├── VeRi │ ├── deit_transreid.yml │ ├── deit_transreid_stride.yml │ ├── vit_base.yml │ ├── vit_transreid.yml │ └── vit_transreid_stride.yml ├── VehicleID │ ├── deit_transreid.yml │ ├── deit_transreid_stride.yml │ ├── vit_base.yml │ ├── vit_transreid.yml │ └── vit_transreid_stride.yml └── transformer_base.yml ├── datasets ├── __init__.py ├── bases.py ├── cuhk03_np_detected.py ├── cuhk03_np_labeled.py ├── dukemtmcreid.py ├── keypoint_test.txt ├── keypoint_train.txt ├── make_dataloader.py ├── market1501.py ├── msmt17.py ├── occ_duke.py ├── preprocessing.py ├── sampler.py ├── sampler_ddp.py ├── vehicleid.py └── veri.py ├── dist_train.sh ├── figs ├── ablation.png ├── framework.png └── sota.png ├── loss ├── __init__.py ├── arcface.py ├── center_loss.py ├── make_loss.py ├── metric_learning.py ├── softmax_loss.py └── triplet_loss.py ├── model ├── __init__.py ├── backbones │ ├── __init__.py │ ├── resnet.py │ └── vit_pytorch.py └── make_model.py ├── processor ├── __init__.py └── processor.py ├── requirements.txt ├── solver ├── __init__.py ├── cosine_lr.py ├── lr_scheduler.py ├── make_optimizer.py ├── scheduler.py └── scheduler_factory.py ├── test.py ├── train.py └── utils ├── __init__.py ├── iotools.py ├── logger.py ├── meter.py ├── metrics.py └── reranking.py /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | # Editor-based HTTP Client requests 5 | /httpRequests/ 6 | # Datasource local storage ignored files 7 | /dataSources/ 8 | /dataSources.local.xml 9 | -------------------------------------------------------------------------------- /.idea/TransReID-main.iml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 12 | -------------------------------------------------------------------------------- /.idea/deployment.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/Project_Default.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 21 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/profiles_settings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 6 | -------------------------------------------------------------------------------- /.idea/misc.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /.idea/modules.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.idea/remote-mappings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 heshuting555 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | # [CVPR2023] PHA: Patch-wise High-frequency Augmentation for Transformer-based Person Re-identification [[pdf]](https://openaccess.thecvf.com/content/CVPR2023/papers/Zhang_PHA_Patch-Wise_High-Frequency_Augmentation_for_Transformer-Based_Person_Re-Identification_CVPR_2023_paper.pdf) 4 | 5 | Official Code for the CVPR 2023 paper [PHA: Patch-wise High-frequency Augmentation for Transformer-based Person Re-identification]. 6 | 7 | 8 | 9 | 10 | ## Requirements 11 | 12 | ### Installation 13 | 14 | ```bash 15 | pip install -r requirements.txt 16 | (we use 32G V100 for training and evaluation.) 17 | ``` 18 | 19 | 20 | 21 | ### Prepare ViT Pre-trained Models 22 | 23 | You need to download the ImageNet pretrained transformer model : [ViT-Base](https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth), 24 | 25 | ## Training 26 | 27 | We utilize 1 GPU for training. 28 | 29 | ```bash 30 | CUDA_VISIBLE_DEVICES=0 python train.py --config_file configs/Cuhk03_labeled/vit_transreid_stride.yml 31 | CUDA_VISIBLE_DEVICES=0 python train.py --config_file configs/Market/vit_transreid_stride.yml 32 | CUDA_VISIBLE_DEVICES=0 python train.py --config_file configs/MSMT17/vit_transreid_stride.yml 33 | ``` 34 | 35 | 36 | ## Citation 37 | 38 | If you find this code useful for your research, please cite our paper 39 | 40 | ``` 41 | @InProceedings{Zhang_2023_CVPR, 42 | author = {Guiwei Zhang, Yongfei Zhang, Tianyu Zhang, Bo Li1, Shiliang Pu}, 43 | title = {PHA: Patch-wise High-frequency Augmentation for Transformer-based Person Re-identification}, 44 | booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 45 | month = {June}, 46 | year = {2023}, 47 | pages = {14133-14142} 48 | } 49 | ``` 50 | ## Acknowledgement 51 | 52 | Our code is based on TransReID. Thanks for the great work! 53 | ```bibtex 54 | @InProceedings{He_2021_ICCV, 55 | author = {He, Shuting and Luo, Hao and Wang, Pichao and Wang, Fan and Li, Hao and Jiang, Wei}, 56 | title = {TransReID: Transformer-Based Object Re-Identification}, 57 | booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, 58 | month = {October}, 59 | year = {2021}, 60 | pages = {15013-15022} 61 | } 62 | ``` -------------------------------------------------------------------------------- /config/__init__.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: sherlock 4 | @contact: sherlockliao01@gmail.com 5 | """ 6 | 7 | from .defaults import _C as cfg 8 | from .defaults import _C as cfg_test 9 | -------------------------------------------------------------------------------- /config/defaults.py: -------------------------------------------------------------------------------- 1 | from yacs.config import CfgNode as CN 2 | 3 | # ----------------------------------------------------------------------------- 4 | # Convention about Training / Test specific parameters 5 | # ----------------------------------------------------------------------------- 6 | # Whenever an argument can be either used for training or for testing, the 7 | # corresponding name will be post-fixed by a _TRAIN for a training parameter, 8 | 9 | # ----------------------------------------------------------------------------- 10 | # Config definition 11 | # ----------------------------------------------------------------------------- 12 | 13 | _C = CN() 14 | # ----------------------------------------------------------------------------- 15 | # MODEL 16 | # ----------------------------------------------------------------------------- 17 | _C.MODEL = CN() 18 | # Using cuda or cpu for training 19 | _C.MODEL.DEVICE = "cuda" 20 | # ID number of GPU 21 | _C.MODEL.DEVICE_ID = '0' 22 | # Name of backbone 23 | _C.MODEL.NAME = 'resnet50' 24 | # Last stride of backbone 25 | _C.MODEL.LAST_STRIDE = 1 26 | # Path to pretrained model of backbone 27 | _C.MODEL.PRETRAIN_PATH = '' 28 | 29 | # Use ImageNet pretrained model to initialize backbone or use self trained model to initialize the whole model 30 | # Options: 'imagenet' , 'self' , 'finetune' 31 | _C.MODEL.PRETRAIN_CHOICE = 'imagenet' 32 | 33 | # If train with BNNeck, options: 'bnneck' or 'no' 34 | _C.MODEL.NECK = 'bnneck' 35 | # If train loss include center loss, options: 'yes' or 'no'. Loss with center loss has different optimizer configuration 36 | _C.MODEL.IF_WITH_CENTER = 'no' 37 | 38 | _C.MODEL.ID_LOSS_TYPE = 'softmax' 39 | _C.MODEL.ID_LOSS_WEIGHT = 1.0 40 | _C.MODEL.TRIPLET_LOSS_WEIGHT = 1.0 41 | 42 | _C.MODEL.METRIC_LOSS_TYPE = 'triplet' 43 | # If train with multi-gpu ddp mode, options: 'True', 'False' 44 | _C.MODEL.DIST_TRAIN = False 45 | # If train with soft triplet loss, options: 'True', 'False' 46 | _C.MODEL.NO_MARGIN = False 47 | # If train with label smooth, options: 'on', 'off' 48 | _C.MODEL.IF_LABELSMOOTH = 'on' 49 | # If train with arcface loss, options: 'True', 'False' 50 | _C.MODEL.COS_LAYER = False 51 | 52 | # Transformer setting 53 | _C.MODEL.DROP_PATH = 0.1 54 | _C.MODEL.DROP_OUT = 0.0 55 | _C.MODEL.ATT_DROP_RATE = 0.0 56 | _C.MODEL.TRANSFORMER_TYPE = 'None' 57 | _C.MODEL.STRIDE_SIZE = [16, 16] 58 | 59 | # JPM Parameter 60 | _C.MODEL.JPM = False 61 | _C.MODEL.SHIFT_NUM = 5 62 | _C.MODEL.SHUFFLE_GROUP = 2 63 | _C.MODEL.DEVIDE_LENGTH = 4 64 | _C.MODEL.RE_ARRANGE = True 65 | 66 | # SIE Parameter 67 | _C.MODEL.SIE_COE = 3.0 68 | _C.MODEL.SIE_CAMERA = False 69 | _C.MODEL.SIE_VIEW = False 70 | 71 | # ----------------------------------------------------------------------------- 72 | # INPUT 73 | # ----------------------------------------------------------------------------- 74 | _C.INPUT = CN() 75 | # Size of the image during training 76 | _C.INPUT.SIZE_TRAIN = [384, 128] 77 | # Size of the image during test 78 | _C.INPUT.SIZE_TEST = [384, 128] 79 | # Random probability for image horizontal flip 80 | _C.INPUT.PROB = 0.5 81 | # Random probability for random erasing 82 | _C.INPUT.RE_PROB = 0.5 83 | # Values to be used for image normalization 84 | _C.INPUT.PIXEL_MEAN = [0.485, 0.456, 0.406] 85 | # Values to be used for image normalization 86 | _C.INPUT.PIXEL_STD = [0.229, 0.224, 0.225] 87 | # Value of padding size 88 | _C.INPUT.PADDING = 10 89 | 90 | # ----------------------------------------------------------------------------- 91 | # Dataset 92 | # ----------------------------------------------------------------------------- 93 | _C.DATASETS = CN() 94 | # List of the dataset names for training, as present in paths_catalog.py 95 | _C.DATASETS.NAMES = ('market1501') 96 | # Root directory where datasets should be used (and downloaded if not found) 97 | _C.DATASETS.ROOT_DIR = ('../data') 98 | 99 | 100 | # ----------------------------------------------------------------------------- 101 | # DataLoader 102 | # ----------------------------------------------------------------------------- 103 | _C.DATALOADER = CN() 104 | # Number of data loading threads 105 | _C.DATALOADER.NUM_WORKERS = 8 106 | # Sampler for data loading 107 | _C.DATALOADER.SAMPLER = 'softmax' 108 | # Number of instance for one batch 109 | _C.DATALOADER.NUM_INSTANCE = 16 110 | 111 | # ---------------------------------------------------------------------------- # 112 | # Solver 113 | # ---------------------------------------------------------------------------- # 114 | _C.SOLVER = CN() 115 | # Name of optimizer 116 | _C.SOLVER.OPTIMIZER_NAME = "Adam" 117 | # Number of max epoches 118 | _C.SOLVER.MAX_EPOCHS = 100 119 | # Base learning rate 120 | _C.SOLVER.BASE_LR = 3e-4 121 | # Whether using larger learning rate for fc layer 122 | _C.SOLVER.LARGE_FC_LR = False 123 | # Factor of learning bias 124 | _C.SOLVER.BIAS_LR_FACTOR = 1 125 | # Factor of learning bias 126 | _C.SOLVER.SEED = 1234 127 | # Momentum 128 | _C.SOLVER.MOMENTUM = 0.9 129 | # Margin of triplet loss 130 | _C.SOLVER.MARGIN = 0.3 131 | # Learning rate of SGD to learn the centers of center loss 132 | _C.SOLVER.CENTER_LR = 0.5 133 | # Balanced weight of center loss 134 | _C.SOLVER.CENTER_LOSS_WEIGHT = 0.0005 135 | 136 | # Settings of weight decay 137 | _C.SOLVER.WEIGHT_DECAY = 0.0005 138 | _C.SOLVER.WEIGHT_DECAY_BIAS = 0.0005 139 | 140 | # decay rate of learning rate 141 | _C.SOLVER.GAMMA = 0.1 142 | # decay step of learning rate 143 | _C.SOLVER.STEPS = (40, 70) 144 | # warm up factor 145 | _C.SOLVER.WARMUP_FACTOR = 0.01 146 | # warm up epochs 147 | _C.SOLVER.WARMUP_EPOCHS = 5 148 | # method of warm up, option: 'constant','linear' 149 | _C.SOLVER.WARMUP_METHOD = "linear" 150 | 151 | _C.SOLVER.COSINE_MARGIN = 0.5 152 | _C.SOLVER.COSINE_SCALE = 30 153 | 154 | # epoch number of saving checkpoints 155 | _C.SOLVER.CHECKPOINT_PERIOD = 10 156 | # iteration of display training log 157 | _C.SOLVER.LOG_PERIOD = 100 158 | # epoch number of validation 159 | _C.SOLVER.EVAL_PERIOD = 10 160 | # Number of images per batch 161 | # This is global, so if we have 8 GPUs and IMS_PER_BATCH = 128, each GPU will 162 | # contain 16 images per batch 163 | _C.SOLVER.IMS_PER_BATCH = 64 164 | 165 | # ---------------------------------------------------------------------------- # 166 | # TEST 167 | # ---------------------------------------------------------------------------- # 168 | 169 | _C.TEST = CN() 170 | # Number of images per batch during test 171 | _C.TEST.IMS_PER_BATCH = 128 172 | # If test with re-ranking, options: 'True','False' 173 | _C.TEST.RE_RANKING = False 174 | # Path to trained model 175 | _C.TEST.WEIGHT = "" 176 | # Which feature of BNNeck to be used for test, before or after BNNneck, options: 'before' or 'after' 177 | _C.TEST.NECK_FEAT = 'after' 178 | # Whether feature is nomalized before test, if yes, it is equivalent to cosine distance 179 | _C.TEST.FEAT_NORM = 'yes' 180 | 181 | # Name for saving the distmat after testing. 182 | _C.TEST.DIST_MAT = "dist_mat.npy" 183 | # Whether calculate the eval score option: 'True', 'False' 184 | _C.TEST.EVAL = False 185 | # ---------------------------------------------------------------------------- # 186 | # Misc options 187 | # ---------------------------------------------------------------------------- # 188 | # Path to checkpoint and saved log of trained model 189 | _C.OUTPUT_DIR = "" 190 | -------------------------------------------------------------------------------- /configs/Cuhk03_labeled/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '~/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('cuhk03_labeled') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 200 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 32 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 128 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: './logs/cuhk_labeled' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/DukeMTMC/deit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('2') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.8 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('dukemtmc') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/dukemtmc_deit_transreid/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/dukemtmc_deit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('6') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 128] 15 | SIZE_TEST: [256, 128] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('dukemtmc') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.008 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/duke_vit_base' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_jpm.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('1') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | JPM: True 13 | RE_ARRANGE: True 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('dukemtmc') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '../logs/duke_vit_jpm/transformer_120.pth' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/duke_vit_jpm' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_sie.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('2') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('dukemtmc') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '../logs/duke_vit_sie/transformer_120.pth' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/duke_vit_sie' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('dukemtmc') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/duke_vit_transreid/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/duke_vit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_transreid_384.yml: -------------------------------------------------------------------------------- 1 | tMODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [384, 128] 19 | SIZE_TEST: [384, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('dukemtmc') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/duke_vit_transreid_384/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/duke_vit_transreid_384' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('1') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: False 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('dukemtmc') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 48 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/DukeMTMC/vit_transreid_stride_384.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [384, 128] 19 | SIZE_TEST: [384, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('dukemtmc') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/duke_vit_transreid_stride_384/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/duke_vit_transreid_stride_384' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/MSMT17/deit_small.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_small_distilled_patch16_224-649709d9.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('1') 10 | TRANSFORMER_TYPE: 'deit_small_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 128] 15 | SIZE_TEST: [256, 128] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.8 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('msmt17') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.005 35 | IMS_PER_BATCH: 64 36 | LARGE_FC_LR: False 37 | CHECKPOINT_PERIOD: 120 38 | LOG_PERIOD: 50 39 | EVAL_PERIOD: 120 40 | WEIGHT_DECAY: 1e-4 41 | WEIGHT_DECAY_BIAS: 1e-4 42 | BIAS_LR_FACTOR: 2 43 | 44 | TEST: 45 | EVAL: True 46 | IMS_PER_BATCH: 256 47 | RE_RANKING: False 48 | WEIGHT: '' 49 | NECK_FEAT: 'before' 50 | FEAT_NORM: 'yes' 51 | 52 | OUTPUT_DIR: '../logs/msmt17_deit_small_try' 53 | 54 | 55 | -------------------------------------------------------------------------------- /configs/MSMT17/deit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('5') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.8 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('msmt17') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.005 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/msmt17_deit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('4') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 128] 15 | SIZE_TEST: [256, 128] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('msmt17') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.008 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/msmt17_vit_base' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_jpm.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('1') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | JPM: True 13 | RE_ARRANGE: True 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('msmt17') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/msmt17_vit_jpm' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_sie.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('2') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('msmt17') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/msmt17_vit_sie' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_small.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/vit_small_p16_224-15ec54c9.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_small_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 128] 15 | SIZE_TEST: [256, 128] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.8 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('msmt17') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.005 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/msmt17_vit_small' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('msmt17') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/msmt17_vit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_transreid_384.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [384, 128] 19 | SIZE_TEST: [384, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('msmt17') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/msmt17_vit_transreid_384' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('msmt17') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 48 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 60 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/MSMT17/vit_transreid_stride_384.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [384, 128] 19 | SIZE_TEST: [384, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('msmt17') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/msmt17_vit_transreid_stride_384' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/Market/deit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('4') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.8 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('market1501') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/0321_market_deit_transreie/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/0321_market_deit_transreie' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/Market/vit_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('7') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 128] 15 | SIZE_TEST: [256, 128] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('market1501') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.008 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '../logs/0321_market_vit_base/transformer_120.pth' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/0321_market_vit_base' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/Market/vit_jpm.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('1') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | JPM: True 13 | RE_ARRANGE: True 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('market1501') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '../logs/0321_market_vit_jpm/transformer_120.pth' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/0321_market_vit_jpm' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/Market/vit_sie.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('7') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('market1501') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/market_vit_sie' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/Market/vit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('5') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('market1501') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/market_vit_transreid/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/market_vit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/Market/vit_transreid_384.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('5') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [384, 128] 19 | SIZE_TEST: [384, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('market1501') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/market_vit_transreid_384/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/market_vit_transreid_384' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/Market/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: False 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('market1501') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 48 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/Market/vit_transreid_stride_384.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [384, 128] 19 | SIZE_TEST: [384, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('market1501') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/0321_market_vit_transreid_stride_384' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/OCC_Duke/deit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('2') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [11, 11] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.8 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('occ_duke') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/occ_duke_deit_transreid_stride11' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/OCC_Duke/vit_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('5') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 128] 15 | SIZE_TEST: [256, 128] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('occ_duke') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.008 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/occ_duke_vit_base' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/OCC_Duke/vit_jpm.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('1') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | JPM: True 13 | RE_ARRANGE: True 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('occ_duke') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/occ_duke_vit_jpm' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/OCC_Duke/vit_sie.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('2') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | 15 | INPUT: 16 | SIZE_TRAIN: [256, 128] 17 | SIZE_TEST: [256, 128] 18 | PROB: 0.5 # random horizontal flip 19 | RE_PROB: 0.5 # random erasing 20 | PADDING: 10 21 | PIXEL_MEAN: [0.5, 0.5, 0.5] 22 | PIXEL_STD: [0.5, 0.5, 0.5] 23 | 24 | DATASETS: 25 | NAMES: ('occ_duke') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.008 37 | IMS_PER_BATCH: 64 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/occ_duke_vit_sie' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/OCC_Duke/vit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/xxx/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('3') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('occ_duke') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/occ_duke_vit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/OCC_Duke/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_COE: 3.0 14 | JPM: True 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 128] 19 | SIZE_TEST: [256, 128] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('occ_duke') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.008 39 | IMS_PER_BATCH: 48 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/occ_duke_vit_transreid_stride_2' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/VeRi/deit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('4') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_VIEW: True 14 | SIE_COE: 3.0 15 | JPM: True 16 | SHIFT_NUM: 8 17 | RE_ARRANGE: True 18 | 19 | INPUT: 20 | SIZE_TRAIN: [256, 256] 21 | SIZE_TEST: [256, 256] 22 | PROB: 0.5 # random horizontal flip 23 | RE_PROB: 0.8 # random erasing 24 | PADDING: 10 25 | 26 | DATASETS: 27 | NAMES: ('veri') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.01 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/veri_deit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/VeRi/deit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_VIEW: True 14 | SIE_COE: 3.0 15 | JPM: True 16 | SHIFT_NUM: 8 17 | RE_ARRANGE: True 18 | 19 | INPUT: 20 | SIZE_TRAIN: [256, 256] 21 | SIZE_TEST: [256, 256] 22 | PROB: 0.5 # random horizontal flip 23 | RE_PROB: 0.8 # random erasing 24 | PADDING: 10 25 | 26 | DATASETS: 27 | NAMES: ('veri') 28 | ROOT_DIR: ('../../datasets') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.01 39 | IMS_PER_BATCH: 64 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/veri_deit_transreid_stride' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/VeRi/vit_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('4') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 256] 15 | SIZE_TEST: [256, 256] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('veri') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.008 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/veri_vit_base' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/VeRi/vit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('4') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | SIE_CAMERA: True 13 | SIE_VIEW: True 14 | SIE_COE: 3.0 15 | JPM: True 16 | SHIFT_NUM: 8 17 | RE_ARRANGE: True 18 | 19 | INPUT: 20 | SIZE_TRAIN: [256, 256] 21 | SIZE_TEST: [256, 256] 22 | PROB: 0.5 # random horizontal flip 23 | RE_PROB: 0.5 # random erasing 24 | PADDING: 10 25 | PIXEL_MEAN: [0.5, 0.5, 0.5] 26 | PIXEL_STD: [0.5, 0.5, 0.5] 27 | 28 | DATASETS: 29 | NAMES: ('veri') 30 | ROOT_DIR: ('../../data') 31 | 32 | DATALOADER: 33 | SAMPLER: 'softmax_triplet' 34 | NUM_INSTANCE: 4 35 | NUM_WORKERS: 8 36 | 37 | SOLVER: 38 | OPTIMIZER_NAME: 'SGD' 39 | MAX_EPOCHS: 120 40 | BASE_LR: 0.01 41 | IMS_PER_BATCH: 64 42 | WARMUP_METHOD: 'linear' 43 | LARGE_FC_LR: False 44 | CHECKPOINT_PERIOD: 120 45 | LOG_PERIOD: 50 46 | EVAL_PERIOD: 120 47 | WEIGHT_DECAY: 1e-4 48 | WEIGHT_DECAY_BIAS: 1e-4 49 | BIAS_LR_FACTOR: 2 50 | 51 | TEST: 52 | EVAL: True 53 | IMS_PER_BATCH: 256 54 | RE_RANKING: False 55 | WEIGHT: '' 56 | NECK_FEAT: 'before' 57 | FEAT_NORM: 'yes' 58 | 59 | OUTPUT_DIR: '../logs/veri_vit_transreid' 60 | 61 | 62 | -------------------------------------------------------------------------------- /configs/VeRi/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('2') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | SIE_CAMERA: True 13 | SIE_VIEW: True 14 | SIE_COE: 3.0 15 | JPM: True 16 | SHIFT_NUM: 8 17 | RE_ARRANGE: True 18 | 19 | INPUT: 20 | SIZE_TRAIN: [256, 256] 21 | SIZE_TEST: [256, 256] 22 | PROB: 0.5 # random horizontal flip 23 | RE_PROB: 0.5 # random erasing 24 | PADDING: 10 25 | PIXEL_MEAN: [0.5, 0.5, 0.5] 26 | PIXEL_STD: [0.5, 0.5, 0.5] 27 | 28 | DATASETS: 29 | NAMES: ('veri') 30 | ROOT_DIR: ('../../data') 31 | 32 | DATALOADER: 33 | SAMPLER: 'softmax_triplet' 34 | NUM_INSTANCE: 4 35 | NUM_WORKERS: 8 36 | 37 | SOLVER: 38 | OPTIMIZER_NAME: 'SGD' 39 | MAX_EPOCHS: 120 40 | BASE_LR: 0.01 41 | IMS_PER_BATCH: 64 42 | WARMUP_METHOD: 'linear' 43 | LARGE_FC_LR: False 44 | CHECKPOINT_PERIOD: 120 45 | LOG_PERIOD: 50 46 | EVAL_PERIOD: 120 47 | WEIGHT_DECAY: 1e-4 48 | WEIGHT_DECAY_BIAS: 1e-4 49 | BIAS_LR_FACTOR: 2 50 | 51 | TEST: 52 | EVAL: True 53 | IMS_PER_BATCH: 256 54 | RE_RANKING: False 55 | WEIGHT: '' 56 | NECK_FEAT: 'before' 57 | FEAT_NORM: 'yes' 58 | 59 | OUTPUT_DIR: '../logs/veri_vit_transreid_stride' 60 | 61 | 62 | -------------------------------------------------------------------------------- /configs/VehicleID/deit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | # DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | DIST_TRAIN: True 13 | JPM: True 14 | SHIFT_NUM: 8 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 256] 19 | SIZE_TEST: [256, 256] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.8 # random erasing 22 | PADDING: 10 23 | 24 | DATASETS: 25 | NAMES: ('VehicleID') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.03 37 | IMS_PER_BATCH: 256 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/vehicleID_deit_transreid' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/VehicleID/deit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/deit_base_distilled_patch16_224-df68dfff.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | # DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | DIST_TRAIN: True 13 | JPM: True 14 | SHIFT_NUM: 8 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 256] 19 | SIZE_TEST: [256, 256] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.8 # random erasing 22 | PADDING: 10 23 | 24 | DATASETS: 25 | NAMES: ('VehicleID') 26 | ROOT_DIR: ('../../data') 27 | 28 | DATALOADER: 29 | SAMPLER: 'softmax_triplet' 30 | NUM_INSTANCE: 4 31 | NUM_WORKERS: 8 32 | 33 | SOLVER: 34 | OPTIMIZER_NAME: 'SGD' 35 | MAX_EPOCHS: 120 36 | BASE_LR: 0.03 37 | IMS_PER_BATCH: 256 38 | WARMUP_METHOD: 'linear' 39 | LARGE_FC_LR: False 40 | CHECKPOINT_PERIOD: 120 41 | LOG_PERIOD: 50 42 | EVAL_PERIOD: 120 43 | WEIGHT_DECAY: 1e-4 44 | WEIGHT_DECAY_BIAS: 1e-4 45 | BIAS_LR_FACTOR: 2 46 | 47 | TEST: 48 | EVAL: True 49 | IMS_PER_BATCH: 256 50 | RE_RANKING: False 51 | WEIGHT: '' 52 | NECK_FEAT: 'before' 53 | FEAT_NORM: 'yes' 54 | 55 | OUTPUT_DIR: '../logs/vehicleID_deit_transreid_stride' 56 | 57 | 58 | -------------------------------------------------------------------------------- /configs/VehicleID/vit_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | # DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 256] 15 | SIZE_TEST: [256, 256] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('VehicleID') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.04 35 | IMS_PER_BATCH: 224 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '../logs/vehicleID_vit_base/transformer_120.pth' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/vehicleID_vit_base' 54 | 55 | 56 | -------------------------------------------------------------------------------- /configs/VehicleID/vit_transreid.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | # DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | # DIST_TRAIN: True 13 | JPM: True 14 | SHIFT_NUM: 8 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 256] 19 | SIZE_TEST: [256, 256] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('VehicleID') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.045 39 | IMS_PER_BATCH: 224 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/vehicleID_vit_transreid/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/vehicleID_vit_transreid' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/VehicleID/vit_transreid_stride.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | # DEVICE_ID: ('0') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [12, 12] 12 | # DIST_TRAIN: True 13 | JPM: True 14 | SHIFT_NUM: 8 15 | RE_ARRANGE: True 16 | 17 | INPUT: 18 | SIZE_TRAIN: [256, 256] 19 | SIZE_TEST: [256, 256] 20 | PROB: 0.5 # random horizontal flip 21 | RE_PROB: 0.5 # random erasing 22 | PADDING: 10 23 | PIXEL_MEAN: [0.5, 0.5, 0.5] 24 | PIXEL_STD: [0.5, 0.5, 0.5] 25 | 26 | DATASETS: 27 | NAMES: ('VehicleID') 28 | ROOT_DIR: ('../../data') 29 | 30 | DATALOADER: 31 | SAMPLER: 'softmax_triplet' 32 | NUM_INSTANCE: 4 33 | NUM_WORKERS: 8 34 | 35 | SOLVER: 36 | OPTIMIZER_NAME: 'SGD' 37 | MAX_EPOCHS: 120 38 | BASE_LR: 0.045 39 | IMS_PER_BATCH: 256 40 | WARMUP_METHOD: 'linear' 41 | LARGE_FC_LR: False 42 | CHECKPOINT_PERIOD: 120 43 | LOG_PERIOD: 50 44 | EVAL_PERIOD: 120 45 | WEIGHT_DECAY: 1e-4 46 | WEIGHT_DECAY_BIAS: 1e-4 47 | BIAS_LR_FACTOR: 2 48 | 49 | TEST: 50 | EVAL: True 51 | IMS_PER_BATCH: 256 52 | RE_RANKING: False 53 | WEIGHT: '../logs/vehicleID_vit_transreid_stride/transformer_120.pth' 54 | NECK_FEAT: 'before' 55 | FEAT_NORM: 'yes' 56 | 57 | OUTPUT_DIR: '../logs/vehicleID_vit_transreid_stride' 58 | 59 | 60 | -------------------------------------------------------------------------------- /configs/transformer_base.yml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | PRETRAIN_CHOICE: 'imagenet' 3 | PRETRAIN_PATH: '/home/heshuting/.cache/torch/checkpoints/jx_vit_base_p16_224-80ecf9dd.pth' 4 | METRIC_LOSS_TYPE: 'triplet' 5 | IF_LABELSMOOTH: 'off' 6 | IF_WITH_CENTER: 'no' 7 | NAME: 'transformer' 8 | NO_MARGIN: True 9 | DEVICE_ID: ('7') 10 | TRANSFORMER_TYPE: 'vit_base_patch16_224_TransReID' 11 | STRIDE_SIZE: [16, 16] 12 | 13 | INPUT: 14 | SIZE_TRAIN: [256, 256] 15 | SIZE_TEST: [256, 256] 16 | PROB: 0.5 # random horizontal flip 17 | RE_PROB: 0.5 # random erasing 18 | PADDING: 10 19 | PIXEL_MEAN: [0.5, 0.5, 0.5] 20 | PIXEL_STD: [0.5, 0.5, 0.5] 21 | 22 | DATASETS: 23 | NAMES: ('dukemtmc') 24 | ROOT_DIR: ('../../data') 25 | 26 | DATALOADER: 27 | SAMPLER: 'softmax_triplet' 28 | NUM_INSTANCE: 4 29 | NUM_WORKERS: 8 30 | 31 | SOLVER: 32 | OPTIMIZER_NAME: 'SGD' 33 | MAX_EPOCHS: 120 34 | BASE_LR: 0.008 35 | IMS_PER_BATCH: 64 36 | WARMUP_METHOD: 'linear' 37 | LARGE_FC_LR: False 38 | CHECKPOINT_PERIOD: 120 39 | LOG_PERIOD: 50 40 | EVAL_PERIOD: 120 41 | WEIGHT_DECAY: 1e-4 42 | WEIGHT_DECAY_BIAS: 1e-4 43 | BIAS_LR_FACTOR: 2 44 | 45 | TEST: 46 | EVAL: True 47 | IMS_PER_BATCH: 256 48 | RE_RANKING: False 49 | WEIGHT: '' 50 | NECK_FEAT: 'before' 51 | FEAT_NORM: 'yes' 52 | 53 | OUTPUT_DIR: '../logs/' 54 | 55 | 56 | -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .make_dataloader import make_dataloader -------------------------------------------------------------------------------- /datasets/bases.py: -------------------------------------------------------------------------------- 1 | from PIL import Image, ImageFile 2 | 3 | from torch.utils.data import Dataset 4 | import os.path as osp 5 | import random 6 | import torch 7 | ImageFile.LOAD_TRUNCATED_IMAGES = True 8 | 9 | 10 | def read_image(img_path): 11 | """Keep reading image until succeed. 12 | This can avoid IOError incurred by heavy IO process.""" 13 | got_img = False 14 | if not osp.exists(img_path): 15 | raise IOError("{} does not exist".format(img_path)) 16 | while not got_img: 17 | try: 18 | img = Image.open(img_path).convert('RGB') 19 | got_img = True 20 | except IOError: 21 | print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path)) 22 | pass 23 | return img 24 | 25 | 26 | class BaseDataset(object): 27 | """ 28 | Base class of reid dataset 29 | """ 30 | 31 | def get_imagedata_info(self, data): 32 | pids, cams, tracks = [], [], [] 33 | 34 | for _, pid, camid, trackid in data: 35 | pids += [pid] 36 | cams += [camid] 37 | tracks += [trackid] 38 | pids = set(pids) 39 | cams = set(cams) 40 | tracks = set(tracks) 41 | num_pids = len(pids) 42 | num_cams = len(cams) 43 | num_imgs = len(data) 44 | num_views = len(tracks) 45 | return num_pids, num_imgs, num_cams, num_views 46 | 47 | def print_dataset_statistics(self): 48 | raise NotImplementedError 49 | 50 | 51 | class BaseImageDataset(BaseDataset): 52 | """ 53 | Base class of image reid dataset 54 | """ 55 | 56 | def print_dataset_statistics(self, train, query, gallery): 57 | num_train_pids, num_train_imgs, num_train_cams, num_train_views = self.get_imagedata_info(train) 58 | num_query_pids, num_query_imgs, num_query_cams, num_train_views = self.get_imagedata_info(query) 59 | num_gallery_pids, num_gallery_imgs, num_gallery_cams, num_train_views = self.get_imagedata_info(gallery) 60 | 61 | print("Dataset statistics:") 62 | print(" ----------------------------------------") 63 | print(" subset | # ids | # images | # cameras") 64 | print(" ----------------------------------------") 65 | print(" train | {:5d} | {:8d} | {:9d}".format(num_train_pids, num_train_imgs, num_train_cams)) 66 | print(" query | {:5d} | {:8d} | {:9d}".format(num_query_pids, num_query_imgs, num_query_cams)) 67 | print(" gallery | {:5d} | {:8d} | {:9d}".format(num_gallery_pids, num_gallery_imgs, num_gallery_cams)) 68 | print(" ----------------------------------------") 69 | 70 | 71 | class ImageDataset(Dataset): 72 | def __init__(self, dataset, transform=None): 73 | self.dataset = dataset 74 | self.transform = transform 75 | 76 | def __len__(self): 77 | return len(self.dataset) 78 | 79 | def __getitem__(self, index): 80 | img_path, pid, camid, trackid = self.dataset[index] 81 | img = read_image(img_path) 82 | 83 | if self.transform is not None: 84 | img = self.transform(img) 85 | 86 | return img, pid, camid, trackid,img_path 87 | -------------------------------------------------------------------------------- /datasets/cuhk03_np_detected.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | 3 | 4 | import glob 5 | import re 6 | 7 | import os.path as osp 8 | import os 9 | from PIL import Image 10 | from .bases import BaseImageDataset 11 | 12 | import numpy as np 13 | 14 | 15 | class CUHK03_NP_detected(BaseImageDataset): 16 | """ 17 | Market1501 18 | Reference: 19 | Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015. 20 | URL: http://www.liangzheng.org/Project/project_reid.html 21 | 22 | Dataset statistics: 23 | # identities: 1501 (+1 for background) 24 | # images: 12936 (train) + 3368 (query) + 15913 (gallery) 25 | """ 26 | dataset_dir = 'cuhk03-np/detected' 27 | 28 | def __init__(self, root='/data2/kzhu/', pseudo_label_subdir='train_mask_annotations', part_num=7, verbose=True, **kwargs): 29 | super(CUHK03_NP_detected, self).__init__() 30 | self.dataset_dir = osp.join(root, self.dataset_dir) 31 | self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') 32 | self.query_dir = osp.join(self.dataset_dir, 'query') 33 | self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') 34 | self.pseudo_label_subdir = pseudo_label_subdir 35 | self.pseudo_label_dir = osp.join(self.dataset_dir, self.pseudo_label_subdir) 36 | self.part_num = part_num 37 | if not osp.exists(self.pseudo_label_dir): 38 | os.mkdir(self.pseudo_label_dir) 39 | 40 | self._check_before_run() 41 | 42 | train = self._process_train_dir(self.train_dir, relabel=True) 43 | query = self._process_test_dir(self.query_dir, relabel=False) 44 | gallery = self._process_test_dir(self.gallery_dir, relabel=False) 45 | 46 | if verbose: 47 | print("=> cuhk03_np_detected loaded") 48 | self.print_dataset_statistics(train, query, gallery) 49 | 50 | self.train = train 51 | self.query = query 52 | self.gallery = gallery 53 | 54 | self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train) 55 | self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query) 56 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery) 57 | 58 | def _check_before_run(self): 59 | """Check if all files are available before going deeper""" 60 | if not osp.exists(self.dataset_dir): 61 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 62 | if not osp.exists(self.train_dir): 63 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 64 | if not osp.exists(self.query_dir): 65 | raise RuntimeError("'{}' is not available".format(self.query_dir)) 66 | if not osp.exists(self.gallery_dir): 67 | raise RuntimeError("'{}' is not available".format(self.gallery_dir)) 68 | 69 | def _process_train_dir(self, dir_path, relabel=False): 70 | img_paths = glob.glob(osp.join(dir_path, '*.png')) 71 | pattern = re.compile(r'([-\d]+)_c(\d)') 72 | 73 | pid_container = set() 74 | for img_path in img_paths: 75 | pid, _ = map(int, pattern.search(img_path).groups()) 76 | if pid == -1: continue # junk images are just ignored 77 | pid_container.add(pid) 78 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 79 | 80 | dataset = [] 81 | for img_path in img_paths: 82 | pid, camid = map(int, pattern.search(img_path).groups()) 83 | if pid == -1: continue # junk images are just ignored 84 | camid -= 1 # index starts from 0 85 | if relabel: pid = pid2label[pid] 86 | pseudo_labels_path=osp.splitext(osp.join(self.pseudo_label_dir, osp.basename(img_path)))[0]+'.png' 87 | dataset.append((img_path, pid, camid, pseudo_labels_path)) 88 | return dataset 89 | 90 | def _process_test_dir(self, dir_path, relabel=False): 91 | img_paths = glob.glob(osp.join(dir_path, '*.png')) 92 | pattern = re.compile(r'([-\d]+)_c(\d)') 93 | 94 | pid_container = set() 95 | for img_path in img_paths: 96 | pid, _ = map(int, pattern.search(img_path).groups()) 97 | pid_container.add(pid) 98 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 99 | 100 | dataset = [] 101 | for img_path in img_paths: 102 | pid, camid = map(int, pattern.search(img_path).groups()) 103 | #assert 1 <= camid <= 8 104 | camid -= 1 # index starts from 0 105 | if relabel: pid = pid2label[pid] 106 | dataset.append((img_path, pid, camid, '')) 107 | 108 | return dataset 109 | -------------------------------------------------------------------------------- /datasets/cuhk03_np_labeled.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | 3 | 4 | import glob 5 | import re 6 | 7 | import os.path as osp 8 | import os 9 | from PIL import Image 10 | from .bases import BaseImageDataset 11 | 12 | import numpy as np 13 | 14 | 15 | class CUHK03_NP_labeled(BaseImageDataset): 16 | """ 17 | Market1501 18 | Reference: 19 | Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015. 20 | URL: http://www.liangzheng.org/Project/project_reid.html 21 | 22 | Dataset statistics: 23 | # identities: 1501 (+1 for background) 24 | # images: 12936 (train) + 3368 (query) + 15913 (gallery) 25 | """ 26 | dataset_dir = 'cuhk03-np/labeled' 27 | 28 | def __init__(self, root='/data2/kzhu/', pseudo_label_subdir='train_mask_annotations', part_num=7, verbose=True, **kwargs): 29 | super(CUHK03_NP_labeled, self).__init__() 30 | self.dataset_dir = osp.join(root, self.dataset_dir) 31 | self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') 32 | self.query_dir = osp.join(self.dataset_dir, 'query') 33 | self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') 34 | self.pseudo_label_subdir = pseudo_label_subdir 35 | self.pseudo_label_dir = osp.join(self.dataset_dir, self.pseudo_label_subdir) 36 | self.part_num = part_num 37 | if not osp.exists(self.pseudo_label_dir): 38 | os.mkdir(self.pseudo_label_dir) 39 | self._check_before_run() 40 | 41 | train = self._process_train_dir(self.train_dir, relabel=True) 42 | query = self._process_test_dir(self.query_dir, relabel=False) 43 | gallery = self._process_test_dir(self.gallery_dir, relabel=False) 44 | 45 | if verbose: 46 | print("=> cuhk03_np_labeled loaded") 47 | self.print_dataset_statistics(train, query, gallery) 48 | 49 | self.train = train 50 | self.query = query 51 | self.gallery = gallery 52 | 53 | self.num_train_pids, self.num_train_imgs, self.num_train_cams,self.num_train_vids = self.get_imagedata_info(self.train) 54 | self.num_query_pids, self.num_query_imgs, self.num_query_cams,self.num_query_vids = self.get_imagedata_info(self.query) 55 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams,self.num_gallery_vids = self.get_imagedata_info(self.gallery) 56 | 57 | def _check_before_run(self): 58 | """Check if all files are available before going deeper""" 59 | if not osp.exists(self.dataset_dir): 60 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 61 | if not osp.exists(self.train_dir): 62 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 63 | if not osp.exists(self.query_dir): 64 | raise RuntimeError("'{}' is not available".format(self.query_dir)) 65 | if not osp.exists(self.gallery_dir): 66 | raise RuntimeError("'{}' is not available".format(self.gallery_dir)) 67 | 68 | def _process_train_dir(self, dir_path, relabel=False): 69 | img_paths = glob.glob(osp.join(dir_path, '*.png')) 70 | pattern = re.compile(r'([-\d]+)_c(\d)') 71 | 72 | pid_container = set() 73 | for img_path in img_paths: 74 | pid, _ = map(int, pattern.search(img_path).groups()) 75 | if pid == -1: continue # junk images are just ignored 76 | pid_container.add(pid) 77 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 78 | 79 | dataset = [] 80 | for img_path in img_paths: 81 | pid, camid = map(int, pattern.search(img_path).groups()) 82 | if pid == -1: continue # junk images are just ignored 83 | camid -= 1 # index starts from 0 84 | if relabel: pid = pid2label[pid] 85 | pseudo_labels_path=osp.splitext(osp.join(self.pseudo_label_dir, osp.basename(img_path)))[0]+'.png' 86 | dataset.append((img_path, pid, camid, 1)) 87 | return dataset 88 | 89 | def _process_test_dir(self, dir_path, relabel=False): 90 | img_paths = glob.glob(osp.join(dir_path, '*.png')) 91 | pattern = re.compile(r'([-\d]+)_c(\d)') 92 | 93 | pid_container = set() 94 | for img_path in img_paths: 95 | pid, _ = map(int, pattern.search(img_path).groups()) 96 | pid_container.add(pid) 97 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 98 | 99 | dataset = [] 100 | for img_path in img_paths: 101 | pid, camid = map(int, pattern.search(img_path).groups()) 102 | #assert 1 <= camid <= 8 103 | camid -= 1 # index starts from 0 104 | if relabel: pid = pid2label[pid] 105 | dataset.append((img_path, pid, camid, 1)) 106 | 107 | return dataset 108 | -------------------------------------------------------------------------------- /datasets/dukemtmcreid.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: liaoxingyu 4 | @contact: liaoxingyu2@jd.com 5 | """ 6 | 7 | import glob 8 | import re 9 | import urllib 10 | import zipfile 11 | 12 | import os.path as osp 13 | 14 | from utils.iotools import mkdir_if_missing 15 | from .bases import BaseImageDataset 16 | 17 | 18 | class DukeMTMCreID(BaseImageDataset): 19 | """ 20 | DukeMTMC-reID 21 | Reference: 22 | 1. Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016. 23 | 2. Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017. 24 | URL: https://github.com/layumi/DukeMTMC-reID_evaluation 25 | 26 | Dataset statistics: 27 | # identities: 1404 (train + query) 28 | # images:16522 (train) + 2228 (query) + 17661 (gallery) 29 | # cameras: 8 30 | """ 31 | dataset_dir = 'dukemtmcreid' 32 | 33 | def __init__(self, root='', verbose=True, pid_begin=0, **kwargs): 34 | super(DukeMTMCreID, self).__init__() 35 | self.dataset_dir = osp.join(root, self.dataset_dir) 36 | 37 | self.dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip' 38 | self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') 39 | self.query_dir = osp.join(self.dataset_dir, 'query') 40 | self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') 41 | self.pid_begin = pid_begin 42 | self._download_data() 43 | self._check_before_run() 44 | 45 | train = self._process_dir(self.train_dir, relabel=True) 46 | query = self._process_dir(self.query_dir, relabel=False) 47 | gallery = self._process_dir(self.gallery_dir, relabel=False) 48 | 49 | if verbose: 50 | print("=> DukeMTMC-reID loaded") 51 | self.print_dataset_statistics(train, query, gallery) 52 | 53 | self.train = train 54 | self.query = query 55 | self.gallery = gallery 56 | 57 | self.num_train_pids, self.num_train_imgs, self.num_train_cams, self.num_train_vids = self.get_imagedata_info(self.train) 58 | self.num_query_pids, self.num_query_imgs, self.num_query_cams, self.num_query_vids = self.get_imagedata_info(self.query) 59 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams, self.num_gallery_vids = self.get_imagedata_info(self.gallery) 60 | 61 | def _download_data(self): 62 | if osp.exists(self.dataset_dir): 63 | print("This dataset has been downloaded.") 64 | return 65 | 66 | print("Creating directory {}".format(self.dataset_dir)) 67 | mkdir_if_missing(self.dataset_dir) 68 | fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url)) 69 | 70 | print("Downloading DukeMTMC-reID dataset") 71 | urllib.request.urlretrieve(self.dataset_url, fpath) 72 | 73 | print("Extracting files") 74 | zip_ref = zipfile.ZipFile(fpath, 'r') 75 | zip_ref.extractall(self.dataset_dir) 76 | zip_ref.close() 77 | 78 | def _check_before_run(self): 79 | """Check if all files are available before going deeper""" 80 | if not osp.exists(self.dataset_dir): 81 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 82 | if not osp.exists(self.train_dir): 83 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 84 | if not osp.exists(self.query_dir): 85 | raise RuntimeError("'{}' is not available".format(self.query_dir)) 86 | if not osp.exists(self.gallery_dir): 87 | raise RuntimeError("'{}' is not available".format(self.gallery_dir)) 88 | 89 | def _process_dir(self, dir_path, relabel=False): 90 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 91 | pattern = re.compile(r'([-\d]+)_c(\d)') 92 | 93 | pid_container = set() 94 | for img_path in img_paths: 95 | pid, _ = map(int, pattern.search(img_path).groups()) 96 | pid_container.add(pid) 97 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 98 | 99 | dataset = [] 100 | cam_container = set() 101 | for img_path in img_paths: 102 | pid, camid = map(int, pattern.search(img_path).groups()) 103 | assert 1 <= camid <= 8 104 | camid -= 1 # index starts from 0 105 | if relabel: pid = pid2label[pid] 106 | dataset.append((img_path, self.pid_begin + pid, camid, 1)) 107 | cam_container.add(camid) 108 | print(cam_container, 'cam_container') 109 | return dataset 110 | -------------------------------------------------------------------------------- /datasets/make_dataloader.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision.transforms as T 3 | from torch.utils.data import DataLoader 4 | 5 | from .bases import ImageDataset 6 | from timm.data.random_erasing import RandomErasing 7 | from .sampler import RandomIdentitySampler 8 | from .dukemtmcreid import DukeMTMCreID 9 | from .market1501 import Market1501 10 | from .msmt17 import MSMT17 11 | from .sampler_ddp import RandomIdentitySampler_DDP 12 | import torch.distributed as dist 13 | from .occ_duke import OCC_DukeMTMCreID 14 | from .vehicleid import VehicleID 15 | from .veri import VeRi 16 | 17 | from .cuhk03_np_detected import CUHK03_NP_detected 18 | from.cuhk03_np_labeled import CUHK03_NP_labeled 19 | from .veri import VeRi 20 | __factory = { 21 | 'market1501': Market1501, 22 | 'dukemtmc': DukeMTMCreID, 23 | 'msmt17': MSMT17, 24 | 'occ_duke': OCC_DukeMTMCreID, 25 | 'veri': VeRi, 26 | 'VehicleID': VehicleID, 27 | 'cuhk03_labeled':CUHK03_NP_labeled, 28 | 'cuhk03_detected':CUHK03_NP_detected 29 | } 30 | 31 | 32 | def train_collate_fn(batch): 33 | """ 34 | # collate_fn这个函数的输入就是一个list，list的长度是一个batch size，list中的每个元素都是__getitem__得到的结果 35 | """ 36 | imgs, pids, camids, viewids , _ = zip(*batch) 37 | pids = torch.tensor(pids, dtype=torch.int64) 38 | viewids = torch.tensor(viewids, dtype=torch.int64) 39 | camids = torch.tensor(camids, dtype=torch.int64) 40 | return torch.stack(imgs, dim=0), pids, camids, viewids,_ 41 | 42 | def val_collate_fn(batch): 43 | imgs, pids, camids, viewids, img_paths = zip(*batch) 44 | viewids = torch.tensor(viewids, dtype=torch.int64) 45 | camids_batch = torch.tensor(camids, dtype=torch.int64) 46 | return torch.stack(imgs, dim=0), pids, camids, camids_batch, viewids, img_paths 47 | 48 | def make_dataloader(cfg): 49 | train_transforms = T.Compose([ 50 | T.Resize(cfg.INPUT.SIZE_TRAIN, interpolation=3), 51 | T.RandomHorizontalFlip(p=cfg.INPUT.PROB), 52 | T.Pad(cfg.INPUT.PADDING), 53 | T.RandomCrop(cfg.INPUT.SIZE_TRAIN), 54 | T.ToTensor(), 55 | T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD), 56 | RandomErasing(probability=cfg.INPUT.RE_PROB, mode='pixel', max_count=1, device='cpu'), 57 | # RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN) 58 | ]) 59 | 60 | val_transforms = T.Compose([ 61 | T.Resize(cfg.INPUT.SIZE_TEST), 62 | T.ToTensor(), 63 | T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD) 64 | ]) 65 | 66 | num_workers = cfg.DATALOADER.NUM_WORKERS 67 | 68 | dataset = __factory[cfg.DATASETS.NAMES](root=cfg.DATASETS.ROOT_DIR) 69 | 70 | train_set = ImageDataset(dataset.train, train_transforms) 71 | train_set_normal = ImageDataset(dataset.train, val_transforms) 72 | num_classes = dataset.num_train_pids 73 | cam_num = dataset.num_train_cams 74 | view_num = dataset.num_train_vids 75 | 76 | if 'triplet' in cfg.DATALOADER.SAMPLER: 77 | if cfg.MODEL.DIST_TRAIN: 78 | print('DIST_TRAIN START') 79 | mini_batch_size = cfg.SOLVER.IMS_PER_BATCH // dist.get_world_size() 80 | data_sampler = RandomIdentitySampler_DDP(dataset.train, cfg.SOLVER.IMS_PER_BATCH, cfg.DATALOADER.NUM_INSTANCE) 81 | batch_sampler = torch.utils.data.sampler.BatchSampler(data_sampler, mini_batch_size, True) 82 | train_loader = torch.utils.data.DataLoader( 83 | train_set, 84 | num_workers=num_workers, 85 | batch_sampler=batch_sampler, 86 | collate_fn=train_collate_fn, 87 | pin_memory=True, 88 | ) 89 | else: 90 | train_loader = DataLoader( 91 | train_set, batch_size=cfg.SOLVER.IMS_PER_BATCH, 92 | sampler=RandomIdentitySampler(dataset.train, cfg.SOLVER.IMS_PER_BATCH, cfg.DATALOADER.NUM_INSTANCE), 93 | num_workers=num_workers, collate_fn=train_collate_fn 94 | ) 95 | elif cfg.DATALOADER.SAMPLER == 'softmax': 96 | print('using softmax sampler') 97 | train_loader = DataLoader( 98 | train_set, batch_size=cfg.SOLVER.IMS_PER_BATCH, shuffle=True, num_workers=num_workers, 99 | collate_fn=train_collate_fn 100 | ) 101 | else: 102 | print('unsupported sampler! expected softmax or triplet but got {}'.format(cfg.SAMPLER)) 103 | 104 | val_set = ImageDataset(dataset.query + dataset.gallery, val_transforms) 105 | 106 | val_loader = DataLoader( 107 | val_set, batch_size=cfg.TEST.IMS_PER_BATCH, shuffle=False, num_workers=num_workers, 108 | collate_fn=val_collate_fn 109 | ) 110 | train_loader_normal = DataLoader( 111 | train_set_normal, batch_size=cfg.TEST.IMS_PER_BATCH, shuffle=False, num_workers=num_workers, 112 | collate_fn=val_collate_fn 113 | ) 114 | return train_loader, train_loader_normal, val_loader, len(dataset.query), num_classes, cam_num, view_num 115 | -------------------------------------------------------------------------------- /datasets/market1501.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: sherlock 4 | @contact: sherlockliao01@gmail.com 5 | """ 6 | 7 | import glob 8 | import re 9 | 10 | import os.path as osp 11 | 12 | from .bases import BaseImageDataset 13 | from collections import defaultdict 14 | import pickle 15 | class Market1501(BaseImageDataset): 16 | """ 17 | Market1501 18 | Reference: 19 | Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015. 20 | URL: http://www.liangzheng.org/Project/project_reid.html 21 | 22 | Dataset statistics: 23 | # identities: 1501 (+1 for background) 24 | # images: 12936 (train) + 3368 (query) + 15913 (gallery) 25 | """ 26 | dataset_dir = 'market1501' 27 | 28 | def __init__(self, root='', verbose=True, pid_begin = 0, **kwargs): 29 | super(Market1501, self).__init__() 30 | self.dataset_dir = osp.join(root, self.dataset_dir) 31 | self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') 32 | self.query_dir = osp.join(self.dataset_dir, 'query') 33 | self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') 34 | 35 | self._check_before_run() 36 | self.pid_begin = pid_begin 37 | train = self._process_dir(self.train_dir, relabel=True) 38 | query = self._process_dir(self.query_dir, relabel=False) 39 | gallery = self._process_dir(self.gallery_dir, relabel=False) 40 | 41 | if verbose: 42 | print("=> Market1501 loaded") 43 | self.print_dataset_statistics(train, query, gallery) 44 | 45 | self.train = train 46 | self.query = query 47 | self.gallery = gallery 48 | 49 | self.num_train_pids, self.num_train_imgs, self.num_train_cams, self.num_train_vids = self.get_imagedata_info(self.train) 50 | self.num_query_pids, self.num_query_imgs, self.num_query_cams, self.num_query_vids = self.get_imagedata_info(self.query) 51 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams, self.num_gallery_vids = self.get_imagedata_info(self.gallery) 52 | 53 | def _check_before_run(self): 54 | """Check if all files are available before going deeper""" 55 | if not osp.exists(self.dataset_dir): 56 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 57 | if not osp.exists(self.train_dir): 58 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 59 | if not osp.exists(self.query_dir): 60 | raise RuntimeError("'{}' is not available".format(self.query_dir)) 61 | if not osp.exists(self.gallery_dir): 62 | raise RuntimeError("'{}' is not available".format(self.gallery_dir)) 63 | 64 | def _process_dir(self, dir_path, relabel=False): 65 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 66 | pattern = re.compile(r'([-\d]+)_c(\d)') 67 | 68 | pid_container = set() 69 | for img_path in sorted(img_paths): 70 | pid, _ = map(int, pattern.search(img_path).groups()) 71 | if pid == -1: continue # junk images are just ignored 72 | pid_container.add(pid) 73 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 74 | dataset = [] 75 | for img_path in sorted(img_paths): 76 | pid, camid = map(int, pattern.search(img_path).groups()) 77 | if pid == -1: continue # junk images are just ignored 78 | assert 0 <= pid <= 1501 # pid == 0 means background 79 | assert 1 <= camid <= 6 80 | camid -= 1 # index starts from 0 81 | if relabel: pid = pid2label[pid] 82 | 83 | dataset.append((img_path, self.pid_begin + pid, camid, 1)) 84 | return dataset 85 | -------------------------------------------------------------------------------- /datasets/msmt17.py: -------------------------------------------------------------------------------- 1 | 2 | import glob 3 | import re 4 | 5 | import os.path as osp 6 | 7 | from .bases import BaseImageDataset 8 | 9 | 10 | class MSMT17(BaseImageDataset): 11 | """ 12 | MSMT17 13 | 14 | Reference: 15 | Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018. 16 | 17 | URL: http://www.pkuvmc.com/publications/msmt17.html 18 | 19 | Dataset statistics: 20 | # identities: 4101 21 | # images: 32621 (train) + 11659 (query) + 82161 (gallery) 22 | # cameras: 15 23 | """ 24 | dataset_dir = 'MSMT17' 25 | 26 | def __init__(self, root='', verbose=True, pid_begin=0, **kwargs): 27 | super(MSMT17, self).__init__() 28 | self.pid_begin = pid_begin 29 | self.dataset_dir = osp.join(root, self.dataset_dir) 30 | self.train_dir = osp.join(self.dataset_dir, 'train') 31 | self.test_dir = osp.join(self.dataset_dir, 'test') 32 | self.list_train_path = osp.join(self.dataset_dir, 'list_train.txt') 33 | self.list_val_path = osp.join(self.dataset_dir, 'list_val.txt') 34 | self.list_query_path = osp.join(self.dataset_dir, 'list_query.txt') 35 | self.list_gallery_path = osp.join(self.dataset_dir, 'list_gallery.txt') 36 | 37 | self._check_before_run() 38 | train = self._process_dir(self.train_dir, self.list_train_path) 39 | val = self._process_dir(self.train_dir, self.list_val_path) 40 | train += val 41 | query = self._process_dir(self.test_dir, self.list_query_path) 42 | gallery = self._process_dir(self.test_dir, self.list_gallery_path) 43 | if verbose: 44 | print("=> MSMT17 loaded") 45 | self.print_dataset_statistics(train, query, gallery) 46 | 47 | self.train = train 48 | self.query = query 49 | self.gallery = gallery 50 | 51 | self.num_train_pids, self.num_train_imgs, self.num_train_cams, self.num_train_vids = self.get_imagedata_info(self.train) 52 | self.num_query_pids, self.num_query_imgs, self.num_query_cams, self.num_query_vids = self.get_imagedata_info(self.query) 53 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams, self.num_gallery_vids = self.get_imagedata_info(self.gallery) 54 | def _check_before_run(self): 55 | """Check if all files are available before going deeper""" 56 | if not osp.exists(self.dataset_dir): 57 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 58 | if not osp.exists(self.train_dir): 59 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 60 | if not osp.exists(self.test_dir): 61 | raise RuntimeError("'{}' is not available".format(self.test_dir)) 62 | 63 | def _process_dir(self, dir_path, list_path): 64 | with open(list_path, 'r') as txt: 65 | lines = txt.readlines() 66 | dataset = [] 67 | pid_container = set() 68 | cam_container = set() 69 | for img_idx, img_info in enumerate(lines): 70 | img_path, pid = img_info.split(' ') 71 | pid = int(pid) # no need to relabel 72 | camid = int(img_path.split('_')[2]) 73 | img_path = osp.join(dir_path, img_path) 74 | dataset.append((img_path, self.pid_begin +pid, camid-1, 1)) 75 | pid_container.add(pid) 76 | cam_container.add(camid) 77 | print(cam_container, 'cam_container') 78 | # check if pid starts from 0 and increments with 1 79 | for idx, pid in enumerate(pid_container): 80 | assert idx == pid, "See code comment for explanation" 81 | return dataset 82 | -------------------------------------------------------------------------------- /datasets/occ_duke.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: liaoxingyu 4 | @contact: liaoxingyu2@jd.com 5 | """ 6 | 7 | import glob 8 | import re 9 | import urllib 10 | import zipfile 11 | 12 | import os.path as osp 13 | 14 | from utils.iotools import mkdir_if_missing 15 | from .bases import BaseImageDataset 16 | 17 | 18 | class OCC_DukeMTMCreID(BaseImageDataset): 19 | """ 20 | DukeMTMC-reID 21 | Reference: 22 | 1. Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016. 23 | 2. Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017. 24 | URL: https://github.com/layumi/DukeMTMC-reID_evaluation 25 | 26 | Dataset statistics: 27 | # identities: 1404 (train + query) 28 | # images:16522 (train) + 2228 (query) + 17661 (gallery) 29 | # cameras: 8 30 | """ 31 | dataset_dir = 'Occluded_Duke' 32 | 33 | def __init__(self, root='', verbose=True, pid_begin=0, **kwargs): 34 | super(OCC_DukeMTMCreID, self).__init__() 35 | self.dataset_dir = osp.join(root, self.dataset_dir) 36 | self.dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip' 37 | self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train') 38 | self.query_dir = osp.join(self.dataset_dir, 'query') 39 | self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test') 40 | self.pid_begin = pid_begin 41 | self._download_data() 42 | self._check_before_run() 43 | 44 | train = self._process_dir(self.train_dir, relabel=True) 45 | query = self._process_dir(self.query_dir, relabel=False) 46 | gallery = self._process_dir(self.gallery_dir, relabel=False) 47 | 48 | if verbose: 49 | print("=> DukeMTMC-reID loaded") 50 | self.print_dataset_statistics(train, query, gallery) 51 | 52 | self.train = train 53 | self.query = query 54 | self.gallery = gallery 55 | 56 | self.num_train_pids, self.num_train_imgs, self.num_train_cams, self.num_train_vids = self.get_imagedata_info(self.train) 57 | self.num_query_pids, self.num_query_imgs, self.num_query_cams, self.num_query_vids = self.get_imagedata_info(self.query) 58 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams, self.num_gallery_vids = self.get_imagedata_info(self.gallery) 59 | 60 | def _download_data(self): 61 | if osp.exists(self.dataset_dir): 62 | print("This dataset has been downloaded.") 63 | return 64 | 65 | print("Creating directory {}".format(self.dataset_dir)) 66 | mkdir_if_missing(self.dataset_dir) 67 | fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url)) 68 | 69 | print("Downloading DukeMTMC-reID dataset") 70 | urllib.request.urlretrieve(self.dataset_url, fpath) 71 | 72 | print("Extracting files") 73 | zip_ref = zipfile.ZipFile(fpath, 'r') 74 | zip_ref.extractall(self.dataset_dir) 75 | zip_ref.close() 76 | 77 | def _check_before_run(self): 78 | """Check if all files are available before going deeper""" 79 | if not osp.exists(self.dataset_dir): 80 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 81 | if not osp.exists(self.train_dir): 82 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 83 | if not osp.exists(self.query_dir): 84 | raise RuntimeError("'{}' is not available".format(self.query_dir)) 85 | if not osp.exists(self.gallery_dir): 86 | raise RuntimeError("'{}' is not available".format(self.gallery_dir)) 87 | 88 | def _process_dir(self, dir_path, relabel=False): 89 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 90 | pattern = re.compile(r'([-\d]+)_c(\d)') 91 | 92 | pid_container = set() 93 | for img_path in img_paths: 94 | pid, _ = map(int, pattern.search(img_path).groups()) 95 | pid_container.add(pid) 96 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 97 | 98 | dataset = [] 99 | cam_container = set() 100 | for img_path in img_paths: 101 | pid, camid = map(int, pattern.search(img_path).groups()) 102 | assert 1 <= camid <= 8 103 | camid -= 1 # index starts from 0 104 | if relabel: pid = pid2label[pid] 105 | dataset.append((img_path, self.pid_begin + pid, camid, 1)) 106 | cam_container.add(camid) 107 | print(cam_container, 'cam_container') 108 | return dataset 109 | -------------------------------------------------------------------------------- /datasets/preprocessing.py: -------------------------------------------------------------------------------- 1 | import random 2 | import math 3 | 4 | 5 | class RandomErasing(object): 6 | """ Randomly selects a rectangle region in an image and erases its pixels. 7 | 'Random Erasing Data Augmentation' by Zhong et al. 8 | See https://arxiv.org/pdf/1708.04896.pdf 9 | Args: 10 | probability: The probability that the Random Erasing operation will be performed. 11 | sl: Minimum proportion of erased area against input image. 12 | sh: Maximum proportion of erased area against input image. 13 | r1: Minimum aspect ratio of erased area. 14 | mean: Erasing value. 15 | """ 16 | 17 | def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)): 18 | self.probability = probability 19 | self.mean = mean 20 | self.sl = sl 21 | self.sh = sh 22 | self.r1 = r1 23 | 24 | def __call__(self, img): 25 | 26 | if random.uniform(0, 1) >= self.probability: 27 | return img 28 | 29 | for attempt in range(100): 30 | area = img.size()[1] * img.size()[2] 31 | 32 | target_area = random.uniform(self.sl, self.sh) * area 33 | aspect_ratio = random.uniform(self.r1, 1 / self.r1) 34 | 35 | h = int(round(math.sqrt(target_area * aspect_ratio))) 36 | w = int(round(math.sqrt(target_area / aspect_ratio))) 37 | 38 | if w < img.size()[2] and h < img.size()[1]: 39 | x1 = random.randint(0, img.size()[1] - h) 40 | y1 = random.randint(0, img.size()[2] - w) 41 | if img.size()[0] == 3: 42 | img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] 43 | img[1, x1:x1 + h, y1:y1 + w] = self.mean[1] 44 | img[2, x1:x1 + h, y1:y1 + w] = self.mean[2] 45 | else: 46 | img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] 47 | return img 48 | 49 | return img 50 | 51 | -------------------------------------------------------------------------------- /datasets/sampler.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data.sampler import Sampler 2 | from collections import defaultdict 3 | import copy 4 | import random 5 | import numpy as np 6 | 7 | class RandomIdentitySampler(Sampler): 8 | """ 9 | Randomly sample N identities, then for each identity, 10 | randomly sample K instances, therefore batch size is N*K. 11 | Args: 12 | - data_source (list): list of (img_path, pid, camid). 13 | - num_instances (int): number of instances per identity in a batch. 14 | - batch_size (int): number of examples in a batch. 15 | """ 16 | 17 | def __init__(self, data_source, batch_size, num_instances): 18 | self.data_source = data_source 19 | self.batch_size = batch_size 20 | self.num_instances = num_instances 21 | self.num_pids_per_batch = self.batch_size // self.num_instances 22 | self.index_dic = defaultdict(list) #dict with list value 23 | #{783: [0, 5, 116, 876, 1554, 2041],...,} 24 | for index, (_, pid, _, _) in enumerate(self.data_source): 25 | self.index_dic[pid].append(index) 26 | self.pids = list(self.index_dic.keys()) 27 | 28 | # estimate number of examples in an epoch 29 | self.length = 0 30 | for pid in self.pids: 31 | idxs = self.index_dic[pid] 32 | num = len(idxs) 33 | if num < self.num_instances: 34 | num = self.num_instances 35 | self.length += num - num % self.num_instances 36 | 37 | def __iter__(self): 38 | batch_idxs_dict = defaultdict(list) 39 | 40 | for pid in self.pids: 41 | idxs = copy.deepcopy(self.index_dic[pid]) 42 | if len(idxs) < self.num_instances: 43 | idxs = np.random.choice(idxs, size=self.num_instances, replace=True) 44 | random.shuffle(idxs) 45 | batch_idxs = [] 46 | for idx in idxs: 47 | batch_idxs.append(idx) 48 | if len(batch_idxs) == self.num_instances: 49 | batch_idxs_dict[pid].append(batch_idxs) 50 | batch_idxs = [] 51 | 52 | avai_pids = copy.deepcopy(self.pids) 53 | final_idxs = [] 54 | 55 | while len(avai_pids) >= self.num_pids_per_batch: 56 | selected_pids = random.sample(avai_pids, self.num_pids_per_batch) 57 | for pid in selected_pids: 58 | batch_idxs = batch_idxs_dict[pid].pop(0) 59 | final_idxs.extend(batch_idxs) 60 | if len(batch_idxs_dict[pid]) == 0: 61 | avai_pids.remove(pid) 62 | 63 | return iter(final_idxs) 64 | 65 | def __len__(self): 66 | return self.length 67 | 68 | -------------------------------------------------------------------------------- /datasets/sampler_ddp.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data.sampler import Sampler 2 | from collections import defaultdict 3 | import copy 4 | import random 5 | import numpy as np 6 | import math 7 | import torch.distributed as dist 8 | _LOCAL_PROCESS_GROUP = None 9 | import torch 10 | import pickle 11 | 12 | def _get_global_gloo_group(): 13 | """ 14 | Return a process group based on gloo backend, containing all the ranks 15 | The result is cached. 16 | """ 17 | if dist.get_backend() == "nccl": 18 | return dist.new_group(backend="gloo") 19 | else: 20 | return dist.group.WORLD 21 | 22 | def _serialize_to_tensor(data, group): 23 | backend = dist.get_backend(group) 24 | assert backend in ["gloo", "nccl"] 25 | device = torch.device("cpu" if backend == "gloo" else "cuda") 26 | 27 | buffer = pickle.dumps(data) 28 | if len(buffer) > 1024 ** 3: 29 | print( 30 | "Rank {} trying to all-gather {:.2f} GB of data on device {}".format( 31 | dist.get_rank(), len(buffer) / (1024 ** 3), device 32 | ) 33 | ) 34 | storage = torch.ByteStorage.from_buffer(buffer) 35 | tensor = torch.ByteTensor(storage).to(device=device) 36 | return tensor 37 | 38 | def _pad_to_largest_tensor(tensor, group): 39 | """ 40 | Returns: 41 | list[int]: size of the tensor, on each rank 42 | Tensor: padded tensor that has the max size 43 | """ 44 | world_size = dist.get_world_size(group=group) 45 | assert ( 46 | world_size >= 1 47 | ), "comm.gather/all_gather must be called from ranks within the given group!" 48 | local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device) 49 | size_list = [ 50 | torch.zeros([1], dtype=torch.int64, device=tensor.device) for _ in range(world_size) 51 | ] 52 | dist.all_gather(size_list, local_size, group=group) 53 | size_list = [int(size.item()) for size in size_list] 54 | 55 | max_size = max(size_list) 56 | 57 | # we pad the tensor because torch all_gather does not support 58 | # gathering tensors of different shapes 59 | if local_size != max_size: 60 | padding = torch.zeros((max_size - local_size,), dtype=torch.uint8, device=tensor.device) 61 | tensor = torch.cat((tensor, padding), dim=0) 62 | return size_list, tensor 63 | 64 | def all_gather(data, group=None): 65 | """ 66 | Run all_gather on arbitrary picklable data (not necessarily tensors). 67 | Args: 68 | data: any picklable object 69 | group: a torch process group. By default, will use a group which 70 | contains all ranks on gloo backend. 71 | Returns: 72 | list[data]: list of data gathered from each rank 73 | """ 74 | if dist.get_world_size() == 1: 75 | return [data] 76 | if group is None: 77 | group = _get_global_gloo_group() 78 | if dist.get_world_size(group) == 1: 79 | return [data] 80 | 81 | tensor = _serialize_to_tensor(data, group) 82 | 83 | size_list, tensor = _pad_to_largest_tensor(tensor, group) 84 | max_size = max(size_list) 85 | 86 | # receiving Tensor from all ranks 87 | tensor_list = [ 88 | torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list 89 | ] 90 | dist.all_gather(tensor_list, tensor, group=group) 91 | 92 | data_list = [] 93 | for size, tensor in zip(size_list, tensor_list): 94 | buffer = tensor.cpu().numpy().tobytes()[:size] 95 | data_list.append(pickle.loads(buffer)) 96 | 97 | return data_list 98 | 99 | def shared_random_seed(): 100 | """ 101 | Returns: 102 | int: a random number that is the same across all workers. 103 | If workers need a shared RNG, they can use this shared seed to 104 | create one. 105 | All workers must call this function, otherwise it will deadlock. 106 | """ 107 | ints = np.random.randint(2 ** 31) 108 | all_ints = all_gather(ints) 109 | return all_ints[0] 110 | 111 | class RandomIdentitySampler_DDP(Sampler): 112 | """ 113 | Randomly sample N identities, then for each identity, 114 | randomly sample K instances, therefore batch size is N*K. 115 | Args: 116 | - data_source (list): list of (img_path, pid, camid). 117 | - num_instances (int): number of instances per identity in a batch. 118 | - batch_size (int): number of examples in a batch. 119 | """ 120 | 121 | def __init__(self, data_source, batch_size, num_instances): 122 | self.data_source = data_source 123 | self.batch_size = batch_size 124 | self.world_size = dist.get_world_size() 125 | self.num_instances = num_instances 126 | self.mini_batch_size = self.batch_size // self.world_size 127 | self.num_pids_per_batch = self.mini_batch_size // self.num_instances 128 | self.index_dic = defaultdict(list) 129 | 130 | for index, (_, pid, _, _) in enumerate(self.data_source): 131 | self.index_dic[pid].append(index) 132 | self.pids = list(self.index_dic.keys()) 133 | 134 | # estimate number of examples in an epoch 135 | self.length = 0 136 | for pid in self.pids: 137 | idxs = self.index_dic[pid] 138 | num = len(idxs) 139 | if num < self.num_instances: 140 | num = self.num_instances 141 | self.length += num - num % self.num_instances 142 | 143 | self.rank = dist.get_rank() 144 | #self.world_size = dist.get_world_size() 145 | self.length //= self.world_size 146 | 147 | def __iter__(self): 148 | seed = shared_random_seed() 149 | np.random.seed(seed) 150 | self._seed = int(seed) 151 | final_idxs = self.sample_list() 152 | length = int(math.ceil(len(final_idxs) * 1.0 / self.world_size)) 153 | #final_idxs = final_idxs[self.rank * length:(self.rank + 1) * length] 154 | final_idxs = self.__fetch_current_node_idxs(final_idxs, length) 155 | self.length = len(final_idxs) 156 | return iter(final_idxs) 157 | 158 | 159 | def __fetch_current_node_idxs(self, final_idxs, length): 160 | total_num = len(final_idxs) 161 | block_num = (length // self.mini_batch_size) 162 | index_target = [] 163 | for i in range(0, block_num * self.world_size, self.world_size): 164 | index = range(self.mini_batch_size * self.rank + self.mini_batch_size * i, min(self.mini_batch_size * self.rank + self.mini_batch_size * (i+1), total_num)) 165 | index_target.extend(index) 166 | index_target_npy = np.array(index_target) 167 | final_idxs = list(np.array(final_idxs)[index_target_npy]) 168 | return final_idxs 169 | 170 | 171 | def sample_list(self): 172 | #np.random.seed(self._seed) 173 | avai_pids = copy.deepcopy(self.pids) 174 | batch_idxs_dict = {} 175 | 176 | batch_indices = [] 177 | while len(avai_pids) >= self.num_pids_per_batch: 178 | selected_pids = np.random.choice(avai_pids, self.num_pids_per_batch, replace=False).tolist() 179 | for pid in selected_pids: 180 | if pid not in batch_idxs_dict: 181 | idxs = copy.deepcopy(self.index_dic[pid]) 182 | if len(idxs) < self.num_instances: 183 | idxs = np.random.choice(idxs, size=self.num_instances, replace=True).tolist() 184 | np.random.shuffle(idxs) 185 | batch_idxs_dict[pid] = idxs 186 | 187 | avai_idxs = batch_idxs_dict[pid] 188 | for _ in range(self.num_instances): 189 | batch_indices.append(avai_idxs.pop(0)) 190 | 191 | if len(avai_idxs) < self.num_instances: avai_pids.remove(pid) 192 | 193 | return batch_indices 194 | 195 | def __len__(self): 196 | return self.length 197 | 198 | -------------------------------------------------------------------------------- /datasets/vehicleid.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import random 5 | import os.path as osp 6 | from .bases import BaseImageDataset 7 | from collections import defaultdict 8 | import pickle 9 | 10 | class VehicleID(BaseImageDataset): 11 | """ 12 | VehicleID 13 | Reference: 14 | Deep Relative Distance Learning: Tell the Difference Between Similar Vehicles 15 | 16 | Dataset statistics: 17 | # train_list: 13164 vehicles for model training 18 | # test_list_800: 800 vehicles for model testing(small test set in paper 19 | # test_list_1600: 1600 vehicles for model testing(medium test set in paper 20 | # test_list_2400: 2400 vehicles for model testing(large test set in paper 21 | # test_list_3200: 3200 vehicles for model testing 22 | # test_list_6000: 6000 vehicles for model testing 23 | # test_list_13164: 13164 vehicles for model testing 24 | """ 25 | dataset_dir = 'VehicleID_V1.0' 26 | 27 | def __init__(self, root='', verbose=True, test_size=800, **kwargs): 28 | super(VehicleID, self).__init__() 29 | self.dataset_dir = osp.join(root, self.dataset_dir) 30 | self.img_dir = osp.join(self.dataset_dir, 'image') 31 | self.split_dir = osp.join(self.dataset_dir, 'train_test_split') 32 | self.train_list = osp.join(self.split_dir, 'train_list.txt') 33 | self.test_size = test_size 34 | 35 | if self.test_size == 800: 36 | self.test_list = osp.join(self.split_dir, 'test_list_800.txt') 37 | elif self.test_size == 1600: 38 | self.test_list = osp.join(self.split_dir, 'test_list_1600.txt') 39 | elif self.test_size == 2400: 40 | self.test_list = osp.join(self.split_dir, 'test_list_2400.txt') 41 | 42 | print(self.test_list) 43 | 44 | self.check_before_run() 45 | 46 | train, query, gallery = self.process_split(relabel=True) 47 | self.train = train 48 | self.query = query 49 | self.gallery = gallery 50 | 51 | if verbose: 52 | print('=> VehicleID loaded') 53 | self.print_dataset_statistics(train, query, gallery) 54 | 55 | self.num_train_pids, self.num_train_imgs, self.num_train_cams, self.num_train_vids = self.get_imagedata_info( 56 | self.train) 57 | self.num_query_pids, self.num_query_imgs, self.num_query_cams, self.num_query_vids = self.get_imagedata_info( 58 | self.query) 59 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams, self.num_gallery_vids = self.get_imagedata_info( 60 | self.gallery) 61 | 62 | def check_before_run(self): 63 | """Check if all files are available before going deeper""" 64 | if not osp.exists(self.dataset_dir): 65 | raise RuntimeError('"{}" is not available'.format(self.dataset_dir)) 66 | if not osp.exists(self.split_dir): 67 | raise RuntimeError('"{}" is not available'.format(self.split_dir)) 68 | if not osp.exists(self.train_list): 69 | raise RuntimeError('"{}" is not available'.format(self.train_list)) 70 | if self.test_size not in [800, 1600, 2400]: 71 | raise RuntimeError('"{}" is not available'.format(self.test_size)) 72 | if not osp.exists(self.test_list): 73 | raise RuntimeError('"{}" is not available'.format(self.test_list)) 74 | 75 | def get_pid2label(self, pids): 76 | pid_container = set(pids) 77 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 78 | return pid2label 79 | 80 | 81 | def parse_img_pids(self, nl_pairs, pid2label=None, cam=0): 82 | # il_pair is the pairs of img name and label 83 | output = [] 84 | for info in nl_pairs: 85 | name = info[0] 86 | pid = info[1] 87 | if pid2label is not None: 88 | pid = pid2label[pid] 89 | camid = cam # use 0 or 1 90 | img_path = osp.join(self.img_dir, name+'.jpg') 91 | viewid = 1 92 | output.append((img_path, pid, camid, viewid)) 93 | return output 94 | 95 | def process_split(self, relabel=False): 96 | # read train paths 97 | train_pid_dict = defaultdict(list) 98 | 99 | # 'train_list.txt' format: 100 | # the first number is the number of image 101 | # the second number is the id of vehicle 102 | with open(self.train_list) as f_train: 103 | train_data = f_train.readlines() 104 | for data in train_data: 105 | name, pid = data.strip().split(' ') 106 | 107 | pid = int(pid) 108 | train_pid_dict[pid].append([name, pid]) 109 | train_pids = list(train_pid_dict.keys()) 110 | num_train_pids = len(train_pids) 111 | assert num_train_pids == 13164, 'There should be 13164 vehicles for training,' \ 112 | ' but but got {}, please check the data'\ 113 | .format(num_train_pids) 114 | # print('num of train ids: {}'.format(num_train_pids)) 115 | test_pid_dict = defaultdict(list) 116 | with open(self.test_list) as f_test: 117 | test_data = f_test.readlines() 118 | for data in test_data: 119 | name, pid = data.split(' ') 120 | pid = int(pid) 121 | test_pid_dict[pid].append([name, pid]) 122 | test_pids = list(test_pid_dict.keys()) 123 | num_test_pids = len(test_pids) 124 | assert num_test_pids == self.test_size, 'There should be {} vehicles for testing,' \ 125 | ' but but got {}, please check the data'\ 126 | .format(self.test_size, num_test_pids) 127 | 128 | train_data = [] 129 | query_data = [] 130 | gallery_data = [] 131 | train_pids = sorted(train_pids) 132 | # for train ids, all images are used in the train set. 133 | for pid in train_pids: 134 | imginfo = train_pid_dict[pid] # imginfo include image name and id 135 | train_data.extend(imginfo) 136 | 137 | # for each test id, random choose one image for gallery 138 | # and the other ones for query. 139 | for pid in test_pids: 140 | imginfo = test_pid_dict[pid] 141 | sample = random.choice(imginfo) 142 | imginfo.remove(sample) 143 | query_data.extend(imginfo) 144 | gallery_data.append(sample) 145 | 146 | if relabel: 147 | train_pid2label = self.get_pid2label(train_pids) 148 | else: 149 | train_pid2label = None 150 | 151 | train = self.parse_img_pids(train_data, train_pid2label) 152 | query = self.parse_img_pids(query_data, cam=0) 153 | gallery = self.parse_img_pids(gallery_data, cam=1) 154 | # attach different camera to prevent eval fail 155 | 156 | return train, query, gallery 157 | -------------------------------------------------------------------------------- /datasets/veri.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import re 3 | import os.path as osp 4 | 5 | from .bases import BaseImageDataset 6 | 7 | 8 | class VeRi(BaseImageDataset): 9 | """ 10 | VeRi-776 11 | Reference: 12 | Liu, Xinchen, et al. "Large-scale vehicle re-identification in urban surveillance videos." ICME 2016. 13 | 14 | URL:https://vehiclereid.github.io/VeRi/ 15 | 16 | Dataset statistics: 17 | # identities: 776 18 | # images: 37778 (train) + 1678 (query) + 11579 (gallery) 19 | # cameras: 20 20 | """ 21 | 22 | dataset_dir = 'VeRi' 23 | 24 | def __init__(self, root='', verbose=True, **kwargs): 25 | super(VeRi, self).__init__() 26 | self.dataset_dir = osp.join(root, self.dataset_dir) 27 | self.train_dir = osp.join(self.dataset_dir, 'image_train') 28 | self.query_dir = osp.join(self.dataset_dir, 'image_query') 29 | self.gallery_dir = osp.join(self.dataset_dir, 'image_test') 30 | 31 | self._check_before_run() 32 | 33 | path_train = 'datasets/keypoint_train.txt' 34 | with open(path_train, 'r') as txt: 35 | lines = txt.readlines() 36 | self.image_map_view_train = {} 37 | for img_idx, img_info in enumerate(lines): 38 | content = img_info.split(' ') 39 | viewid = int(content[-1]) 40 | self.image_map_view_train[osp.basename(content[0])] = viewid 41 | 42 | path_test = 'datasets/keypoint_test.txt' 43 | with open(path_test, 'r') as txt: 44 | lines = txt.readlines() 45 | self.image_map_view_test = {} 46 | for img_idx, img_info in enumerate(lines): 47 | content = img_info.split(' ') 48 | viewid = int(content[-1]) 49 | self.image_map_view_test[osp.basename(content[0])] = viewid 50 | 51 | train = self._process_dir(self.train_dir, relabel=True) 52 | query = self._process_dir(self.query_dir, relabel=False) 53 | gallery = self._process_dir(self.gallery_dir, relabel=False) 54 | 55 | if verbose: 56 | print("=> VeRi-776 loaded") 57 | self.print_dataset_statistics(train, query, gallery) 58 | 59 | self.train = train 60 | self.query = query 61 | self.gallery = gallery 62 | 63 | self.num_train_pids, self.num_train_imgs, self.num_train_cams, self.num_train_vids = self.get_imagedata_info( 64 | self.train) 65 | self.num_query_pids, self.num_query_imgs, self.num_query_cams, self.num_query_vids = self.get_imagedata_info( 66 | self.query) 67 | self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams, self.num_gallery_vids = self.get_imagedata_info( 68 | self.gallery) 69 | 70 | def _check_before_run(self): 71 | """Check if all files are available before going deeper""" 72 | if not osp.exists(self.dataset_dir): 73 | raise RuntimeError("'{}' is not available".format(self.dataset_dir)) 74 | if not osp.exists(self.train_dir): 75 | raise RuntimeError("'{}' is not available".format(self.train_dir)) 76 | if not osp.exists(self.query_dir): 77 | raise RuntimeError("'{}' is not available".format(self.query_dir)) 78 | if not osp.exists(self.gallery_dir): 79 | raise RuntimeError("'{}' is not available".format(self.gallery_dir)) 80 | 81 | def _process_dir(self, dir_path, relabel=False): 82 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 83 | pattern = re.compile(r'([-\d]+)_c(\d+)') 84 | 85 | pid_container = set() 86 | for img_path in img_paths: 87 | pid, _ = map(int, pattern.search(img_path).groups()) 88 | if pid == -1: continue # junk images are just ignored 89 | pid_container.add(pid) 90 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 91 | 92 | view_container = set() 93 | dataset = [] 94 | count = 0 95 | for img_path in img_paths: 96 | pid, camid = map(int, pattern.search(img_path).groups()) 97 | if pid == -1: continue # junk images are just ignored 98 | assert 0 <= pid <= 776 # pid == 0 means background 99 | assert 1 <= camid <= 20 100 | camid -= 1 # index starts from 0 101 | if relabel: pid = pid2label[pid] 102 | 103 | if osp.basename(img_path) not in self.image_map_view_train.keys(): 104 | try: 105 | viewid = self.image_map_view_test[osp.basename(img_path)] 106 | except: 107 | count += 1 108 | # print(img_path, 'img_path') 109 | continue 110 | else: 111 | viewid = self.image_map_view_train[osp.basename(img_path)] 112 | view_container.add(viewid) 113 | dataset.append((img_path, pid, camid, viewid)) 114 | print(view_container, 'view_container') 115 | print(count, 'samples without viewpoint annotations') 116 | return dataset 117 | 118 | -------------------------------------------------------------------------------- /dist_train.sh: -------------------------------------------------------------------------------- 1 | # train 2 | CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 --master_port 66666 train.py --config_file configs/VehicleID/vit_base.yml MODEL.DIST_TRAIN True 3 | # test 4 | python test.py --config_file configs/VehicleID/vit_base.yml MODEL.DIST_TRAIN False MODEL.DEVICE_ID "('0')" 5 | -------------------------------------------------------------------------------- /figs/ablation.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhangguiwei610/PHA/fca6ece297036646c99d4af6496df237b94a8606/figs/ablation.png -------------------------------------------------------------------------------- /figs/framework.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhangguiwei610/PHA/fca6ece297036646c99d4af6496df237b94a8606/figs/framework.png -------------------------------------------------------------------------------- /figs/sota.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhangguiwei610/PHA/fca6ece297036646c99d4af6496df237b94a8606/figs/sota.png -------------------------------------------------------------------------------- /loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .make_loss import make_loss 2 | from .arcface import ArcFace -------------------------------------------------------------------------------- /loss/arcface.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn import Parameter 5 | import math 6 | 7 | 8 | class ArcFace(nn.Module): 9 | def __init__(self, in_features, out_features, s=30.0, m=0.50, bias=False): 10 | super(ArcFace, self).__init__() 11 | self.in_features = in_features 12 | self.out_features = out_features 13 | self.s = s 14 | self.m = m 15 | self.cos_m = math.cos(m) 16 | self.sin_m = math.sin(m) 17 | 18 | self.th = math.cos(math.pi - m) 19 | self.mm = math.sin(math.pi - m) * m 20 | 21 | self.weight = Parameter(torch.Tensor(out_features, in_features)) 22 | if bias: 23 | self.bias = Parameter(torch.Tensor(out_features)) 24 | else: 25 | self.register_parameter('bias', None) 26 | self.reset_parameters() 27 | 28 | def reset_parameters(self): 29 | nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 30 | if self.bias is not None: 31 | fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) 32 | bound = 1 / math.sqrt(fan_in) 33 | nn.init.uniform_(self.bias, -bound, bound) 34 | 35 | def forward(self, input, label): 36 | cosine = F.linear(F.normalize(input), F.normalize(self.weight)) 37 | sine = torch.sqrt((1.0 - torch.pow(cosine, 2)).clamp(0, 1)) 38 | phi = cosine * self.cos_m - sine * self.sin_m 39 | phi = torch.where(cosine > self.th, phi, cosine - self.mm) 40 | # --------------------------- convert label to one-hot --------------------------- 41 | # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda') 42 | one_hot = torch.zeros(cosine.size(), device='cuda') 43 | one_hot.scatter_(1, label.view(-1, 1).long(), 1) 44 | # -------------torch.where(out_i = {x_i if condition_i else y_i) ------------- 45 | output = (one_hot * phi) + ( 46 | (1.0 - one_hot) * cosine) # you can use torch.where if your torch.__version__ is 0.4 47 | output *= self.s 48 | # print(output) 49 | 50 | return output 51 | 52 | class CircleLoss(nn.Module): 53 | def __init__(self, in_features, num_classes, s=256, m=0.25): 54 | super(CircleLoss, self).__init__() 55 | self.weight = Parameter(torch.Tensor(num_classes, in_features)) 56 | self.s = s 57 | self.m = m 58 | self._num_classes = num_classes 59 | self.reset_parameters() 60 | 61 | 62 | def reset_parameters(self): 63 | nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 64 | 65 | def __call__(self, bn_feat, targets): 66 | 67 | sim_mat = F.linear(F.normalize(bn_feat), F.normalize(self.weight)) 68 | alpha_p = torch.clamp_min(-sim_mat.detach() + 1 + self.m, min=0.) 69 | alpha_n = torch.clamp_min(sim_mat.detach() + self.m, min=0.) 70 | delta_p = 1 - self.m 71 | delta_n = self.m 72 | 73 | s_p = self.s * alpha_p * (sim_mat - delta_p) 74 | s_n = self.s * alpha_n * (sim_mat - delta_n) 75 | 76 | targets = F.one_hot(targets, num_classes=self._num_classes) 77 | 78 | pred_class_logits = targets * s_p + (1.0 - targets) * s_n 79 | 80 | return pred_class_logits -------------------------------------------------------------------------------- /loss/center_loss.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | import torch 4 | from torch import nn 5 | 6 | 7 | class CenterLoss(nn.Module): 8 | """Center loss. 9 | 10 | Reference: 11 | Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016. 12 | 13 | Args: 14 | num_classes (int): number of classes. 15 | feat_dim (int): feature dimension. 16 | """ 17 | 18 | def __init__(self, num_classes=751, feat_dim=2048, use_gpu=True): 19 | super(CenterLoss, self).__init__() 20 | self.num_classes = num_classes 21 | self.feat_dim = feat_dim 22 | self.use_gpu = use_gpu 23 | 24 | if self.use_gpu: 25 | self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda()) 26 | else: 27 | self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim)) 28 | 29 | def forward(self, x, labels): 30 | """ 31 | Args: 32 | x: feature matrix with shape (batch_size, feat_dim). 33 | labels: ground truth labels with shape (num_classes). 34 | """ 35 | assert x.size(0) == labels.size(0), "features.size(0) is not equal to labels.size(0)" 36 | 37 | batch_size = x.size(0) 38 | distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \ 39 | torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t() 40 | distmat.addmm_(1, -2, x, self.centers.t()) 41 | 42 | classes = torch.arange(self.num_classes).long() 43 | if self.use_gpu: classes = classes.cuda() 44 | labels = labels.unsqueeze(1).expand(batch_size, self.num_classes) 45 | mask = labels.eq(classes.expand(batch_size, self.num_classes)) 46 | 47 | dist = [] 48 | for i in range(batch_size): 49 | value = distmat[i][mask[i]] 50 | value = value.clamp(min=1e-12, max=1e+12) # for numerical stability 51 | dist.append(value) 52 | dist = torch.cat(dist) 53 | loss = dist.mean() 54 | return loss 55 | 56 | 57 | if __name__ == '__main__': 58 | use_gpu = False 59 | center_loss = CenterLoss(use_gpu=use_gpu) 60 | features = torch.rand(16, 2048) 61 | targets = torch.Tensor([0, 1, 2, 3, 2, 3, 1, 4, 5, 3, 2, 1, 0, 0, 5, 4]).long() 62 | if use_gpu: 63 | features = torch.rand(16, 2048).cuda() 64 | targets = torch.Tensor([0, 1, 2, 3, 2, 3, 1, 4, 5, 3, 2, 1, 0, 0, 5, 4]).cuda() 65 | 66 | loss = center_loss(features, targets) 67 | print(loss) 68 | -------------------------------------------------------------------------------- /loss/make_loss.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: liaoxingyu 4 | @contact: sherlockliao01@gmail.com 5 | """ 6 | 7 | import torch.nn.functional as F 8 | from .softmax_loss import CrossEntropyLabelSmooth, LabelSmoothingCrossEntropy 9 | from .triplet_loss import TripletLoss 10 | from .center_loss import CenterLoss 11 | 12 | 13 | def make_loss(cfg, num_classes): # modified by gu 14 | sampler = cfg.DATALOADER.SAMPLER 15 | feat_dim = 2048 16 | center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=True) # center loss 17 | if 'triplet' in cfg.MODEL.METRIC_LOSS_TYPE: 18 | if cfg.MODEL.NO_MARGIN: 19 | triplet = TripletLoss() 20 | print("using soft triplet loss for training") 21 | else: 22 | triplet = TripletLoss(cfg.SOLVER.MARGIN) # triplet loss 23 | print("using triplet loss with margin:{}".format(cfg.SOLVER.MARGIN)) 24 | else: 25 | print('expected METRIC_LOSS_TYPE should be triplet' 26 | 'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE)) 27 | 28 | if cfg.MODEL.IF_LABELSMOOTH == 'on': 29 | xent = CrossEntropyLabelSmooth(num_classes=num_classes) 30 | print("label smooth on, numclasses:", num_classes) 31 | 32 | if sampler == 'softmax': 33 | def loss_func(score, feat, target): 34 | return F.cross_entropy(score, target) 35 | 36 | elif cfg.DATALOADER.SAMPLER == 'softmax_triplet': 37 | def loss_func(score, feat, target, target_cam): 38 | if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet': 39 | if cfg.MODEL.IF_LABELSMOOTH == 'on': 40 | if isinstance(score, list): 41 | ID_LOSS = [xent(scor, target) for scor in score[1:]] 42 | ID_LOSS = sum(ID_LOSS) / len(ID_LOSS) 43 | ID_LOSS = 0.5 * ID_LOSS + 0.5 * xent(score[0], target) 44 | else: 45 | ID_LOSS = xent(score, target) 46 | 47 | if isinstance(feat, list): 48 | TRI_LOSS = [triplet(feats, target)[0] for feats in feat[1:]] 49 | TRI_LOSS = sum(TRI_LOSS) / len(TRI_LOSS) 50 | TRI_LOSS = 0.5 * TRI_LOSS + 0.5 * triplet(feat[0], target)[0] 51 | else: 52 | TRI_LOSS = triplet(feat, target)[0] 53 | 54 | return cfg.MODEL.ID_LOSS_WEIGHT * ID_LOSS + \ 55 | cfg.MODEL.TRIPLET_LOSS_WEIGHT * TRI_LOSS 56 | else: 57 | if isinstance(score, list): 58 | ID_LOSS = [F.cross_entropy(scor, target) for scor in score[1:]] 59 | ID_LOSS = sum(ID_LOSS) / len(ID_LOSS) 60 | ID_LOSS = 0.5 * ID_LOSS + 0.5 * F.cross_entropy(score[0], target) 61 | else: 62 | ID_LOSS = F.cross_entropy(score, target) 63 | 64 | if isinstance(feat, list): 65 | TRI_LOSS = [triplet(feats, target)[0] for feats in feat[1:]] 66 | TRI_LOSS = sum(TRI_LOSS) / len(TRI_LOSS) 67 | TRI_LOSS = 0.5 * TRI_LOSS + 0.5 * triplet(feat[0], target)[0] 68 | else: 69 | TRI_LOSS = triplet(feat, target)[0] 70 | 71 | return cfg.MODEL.ID_LOSS_WEIGHT * ID_LOSS + \ 72 | cfg.MODEL.TRIPLET_LOSS_WEIGHT * TRI_LOSS 73 | else: 74 | print('expected METRIC_LOSS_TYPE should be triplet' 75 | 'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE)) 76 | 77 | else: 78 | print('expected sampler should be softmax, triplet, softmax_triplet or softmax_triplet_center' 79 | 'but got {}'.format(cfg.DATALOADER.SAMPLER)) 80 | return loss_func, center_criterion 81 | 82 | 83 | -------------------------------------------------------------------------------- /loss/metric_learning.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import torch.autograd 5 | from torch.nn import Parameter 6 | import math 7 | 8 | 9 | class ContrastiveLoss(nn.Module): 10 | def __init__(self, margin=0.3, **kwargs): 11 | super(ContrastiveLoss, self).__init__() 12 | self.margin = margin 13 | 14 | def forward(self, inputs, targets): 15 | n = inputs.size(0) 16 | # Compute similarity matrix 17 | sim_mat = torch.matmul(inputs, inputs.t()) 18 | targets = targets 19 | loss = list() 20 | c = 0 21 | 22 | for i in range(n): 23 | pos_pair_ = torch.masked_select(sim_mat[i], targets == targets[i]) 24 | 25 | # move itself 26 | pos_pair_ = torch.masked_select(pos_pair_, pos_pair_ < 1) 27 | neg_pair_ = torch.masked_select(sim_mat[i], targets != targets[i]) 28 | 29 | pos_pair_ = torch.sort(pos_pair_)[0] 30 | neg_pair_ = torch.sort(neg_pair_)[0] 31 | 32 | neg_pair = torch.masked_select(neg_pair_, neg_pair_ > self.margin) 33 | 34 | neg_loss = 0 35 | 36 | pos_loss = torch.sum(-pos_pair_ + 1) 37 | if len(neg_pair) > 0: 38 | neg_loss = torch.sum(neg_pair) 39 | loss.append(pos_loss + neg_loss) 40 | 41 | loss = sum(loss) / n 42 | return loss 43 | 44 | 45 | class CircleLoss(nn.Module): 46 | def __init__(self, in_features, num_classes, s=256, m=0.25): 47 | super(CircleLoss, self).__init__() 48 | self.weight = Parameter(torch.Tensor(num_classes, in_features)) 49 | self.s = s 50 | self.m = m 51 | self._num_classes = num_classes 52 | self.reset_parameters() 53 | 54 | 55 | def reset_parameters(self): 56 | nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 57 | 58 | def __call__(self, bn_feat, targets): 59 | 60 | sim_mat = F.linear(F.normalize(bn_feat), F.normalize(self.weight)) 61 | alpha_p = torch.clamp_min(-sim_mat.detach() + 1 + self.m, min=0.) 62 | alpha_n = torch.clamp_min(sim_mat.detach() + self.m, min=0.) 63 | delta_p = 1 - self.m 64 | delta_n = self.m 65 | 66 | s_p = self.s * alpha_p * (sim_mat - delta_p) 67 | s_n = self.s * alpha_n * (sim_mat - delta_n) 68 | 69 | targets = F.one_hot(targets, num_classes=self._num_classes) 70 | 71 | pred_class_logits = targets * s_p + (1.0 - targets) * s_n 72 | 73 | return pred_class_logits 74 | 75 | 76 | class Arcface(nn.Module): 77 | r"""Implement of large margin arc distance: : 78 | Args: 79 | in_features: size of each input sample 80 | out_features: size of each output sample 81 | s: norm of input feature 82 | m: margin 83 | cos(theta + m) 84 | """ 85 | def __init__(self, in_features, out_features, s=30.0, m=0.30, easy_margin=False, ls_eps=0.0): 86 | super(Arcface, self).__init__() 87 | self.in_features = in_features 88 | self.out_features = out_features 89 | self.s = s 90 | self.m = m 91 | self.ls_eps = ls_eps # label smoothing 92 | self.weight = Parameter(torch.FloatTensor(out_features, in_features)) 93 | nn.init.xavier_uniform_(self.weight) 94 | 95 | self.easy_margin = easy_margin 96 | self.cos_m = math.cos(m) 97 | self.sin_m = math.sin(m) 98 | self.th = math.cos(math.pi - m) 99 | self.mm = math.sin(math.pi - m) * m 100 | 101 | def forward(self, input, label): 102 | # --------------------------- cos(theta) & phi(theta) --------------------------- 103 | cosine = F.linear(F.normalize(input), F.normalize(self.weight)) 104 | sine = torch.sqrt(1.0 - torch.pow(cosine, 2)) 105 | phi = cosine * self.cos_m - sine * self.sin_m 106 | phi = phi.type_as(cosine) 107 | if self.easy_margin: 108 | phi = torch.where(cosine > 0, phi, cosine) 109 | else: 110 | phi = torch.where(cosine > self.th, phi, cosine - self.mm) 111 | # --------------------------- convert label to one-hot --------------------------- 112 | # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda') 113 | one_hot = torch.zeros(cosine.size(), device='cuda') 114 | one_hot.scatter_(1, label.view(-1, 1).long(), 1) 115 | if self.ls_eps > 0: 116 | one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features 117 | # -------------torch.where(out_i = {x_i if condition_i else y_i) ------------- 118 | output = (one_hot * phi) + ((1.0 - one_hot) * cosine) 119 | output *= self.s 120 | 121 | return output 122 | 123 | 124 | class Cosface(nn.Module): 125 | r"""Implement of large margin cosine distance: : 126 | Args: 127 | in_features: size of each input sample 128 | out_features: size of each output sample 129 | s: norm of input feature 130 | m: margin 131 | cos(theta) - m 132 | """ 133 | 134 | def __init__(self, in_features, out_features, s=30.0, m=0.30): 135 | super(Cosface, self).__init__() 136 | self.in_features = in_features 137 | self.out_features = out_features 138 | self.s = s 139 | self.m = m 140 | self.weight = Parameter(torch.FloatTensor(out_features, in_features)) 141 | nn.init.xavier_uniform_(self.weight) 142 | 143 | def forward(self, input, label): 144 | # --------------------------- cos(theta) & phi(theta) --------------------------- 145 | cosine = F.linear(F.normalize(input), F.normalize(self.weight)) 146 | phi = cosine - self.m 147 | # --------------------------- convert label to one-hot --------------------------- 148 | one_hot = torch.zeros(cosine.size(), device='cuda') 149 | # one_hot = one_hot.cuda() if cosine.is_cuda else one_hot 150 | one_hot.scatter_(1, label.view(-1, 1).long(), 1) 151 | # -------------torch.where(out_i = {x_i if condition_i else y_i) ------------- 152 | output = (one_hot * phi) + ((1.0 - one_hot) * cosine) # you can use torch.where if your torch.__version__ is 0.4 153 | output *= self.s 154 | # print(output) 155 | 156 | return output 157 | 158 | def __repr__(self): 159 | return self.__class__.__name__ + '(' \ 160 | + 'in_features=' + str(self.in_features) \ 161 | + ', out_features=' + str(self.out_features) \ 162 | + ', s=' + str(self.s) \ 163 | + ', m=' + str(self.m) + ')' 164 | 165 | 166 | class AMSoftmax(nn.Module): 167 | def __init__(self, in_features, out_features, s=30.0, m=0.30): 168 | super(AMSoftmax, self).__init__() 169 | self.m = m 170 | self.s = s 171 | self.in_feats = in_features 172 | self.W = torch.nn.Parameter(torch.randn(in_features, out_features), requires_grad=True) 173 | self.ce = nn.CrossEntropyLoss() 174 | nn.init.xavier_normal_(self.W, gain=1) 175 | 176 | def forward(self, x, lb): 177 | assert x.size()[0] == lb.size()[0] 178 | assert x.size()[1] == self.in_feats 179 | x_norm = torch.norm(x, p=2, dim=1, keepdim=True).clamp(min=1e-12) 180 | x_norm = torch.div(x, x_norm) 181 | w_norm = torch.norm(self.W, p=2, dim=0, keepdim=True).clamp(min=1e-12) 182 | w_norm = torch.div(self.W, w_norm) 183 | costh = torch.mm(x_norm, w_norm) 184 | # print(x_norm.shape, w_norm.shape, costh.shape) 185 | lb_view = lb.view(-1, 1) 186 | delt_costh = torch.zeros(costh.size(), device='cuda').scatter_(1, lb_view, self.m) 187 | costh_m = costh - delt_costh 188 | costh_m_s = self.s * costh_m 189 | return costh_m_s -------------------------------------------------------------------------------- /loss/softmax_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import functional as F 4 | class CrossEntropyLabelSmooth(nn.Module): 5 | """Cross entropy loss with label smoothing regularizer. 6 | 7 | Reference: 8 | Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016. 9 | Equation: y = (1 - epsilon) * y + epsilon / K. 10 | 11 | Args: 12 | num_classes (int): number of classes. 13 | epsilon (float): weight. 14 | """ 15 | 16 | def __init__(self, num_classes, epsilon=0.1, use_gpu=True): 17 | super(CrossEntropyLabelSmooth, self).__init__() 18 | self.num_classes = num_classes 19 | self.epsilon = epsilon 20 | self.use_gpu = use_gpu 21 | self.logsoftmax = nn.LogSoftmax(dim=1) 22 | 23 | def forward(self, inputs, targets): 24 | """ 25 | Args: 26 | inputs: prediction matrix (before softmax) with shape (batch_size, num_classes) 27 | targets: ground truth labels with shape (num_classes) 28 | """ 29 | log_probs = self.logsoftmax(inputs) 30 | targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1) 31 | if self.use_gpu: targets = targets.cuda() 32 | targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes 33 | loss = (- targets * log_probs).mean(0).sum() 34 | return loss 35 | 36 | class LabelSmoothingCrossEntropy(nn.Module): 37 | """ 38 | NLL loss with label smoothing. 39 | """ 40 | def __init__(self, smoothing=0.1): 41 | """ 42 | Constructor for the LabelSmoothing module. 43 | :param smoothing: label smoothing factor 44 | """ 45 | super(LabelSmoothingCrossEntropy, self).__init__() 46 | assert smoothing < 1.0 47 | self.smoothing = smoothing 48 | self.confidence = 1. - smoothing 49 | 50 | def forward(self, x, target): 51 | logprobs = F.log_softmax(x, dim=-1) 52 | nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1)) 53 | nll_loss = nll_loss.squeeze(1) 54 | smooth_loss = -logprobs.mean(dim=-1) 55 | loss = self.confidence * nll_loss + self.smoothing * smooth_loss 56 | return loss.mean() -------------------------------------------------------------------------------- /loss/triplet_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | 5 | def normalize(x, axis=-1): 6 | """Normalizing to unit length along the specified dimension. 7 | Args: 8 | x: pytorch Variable 9 | Returns: 10 | x: pytorch Variable, same shape as input 11 | """ 12 | x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12) 13 | return x 14 | 15 | 16 | def euclidean_dist(x, y): 17 | """ 18 | Args: 19 | x: pytorch Variable, with shape [m, d] 20 | y: pytorch Variable, with shape [n, d] 21 | Returns: 22 | dist: pytorch Variable, with shape [m, n] 23 | """ 24 | m, n = x.size(0), y.size(0) 25 | xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n) 26 | yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t() 27 | dist = xx + yy 28 | dist = dist - 2 * torch.matmul(x, y.t()) 29 | # dist.addmm_(1, -2, x, y.t()) 30 | dist = dist.clamp(min=1e-12).sqrt() # for numerical stability 31 | return dist 32 | 33 | 34 | def cosine_dist(x, y): 35 | """ 36 | Args: 37 | x: pytorch Variable, with shape [m, d] 38 | y: pytorch Variable, with shape [n, d] 39 | Returns: 40 | dist: pytorch Variable, with shape [m, n] 41 | """ 42 | m, n = x.size(0), y.size(0) 43 | x_norm = torch.pow(x, 2).sum(1, keepdim=True).sqrt().expand(m, n) 44 | y_norm = torch.pow(y, 2).sum(1, keepdim=True).sqrt().expand(n, m).t() 45 | xy_intersection = torch.mm(x, y.t()) 46 | dist = xy_intersection/(x_norm * y_norm) 47 | dist = (1. - dist) / 2 48 | return dist 49 | 50 | 51 | def hard_example_mining(dist_mat, labels, return_inds=False): 52 | """For each anchor, find the hardest positive and negative sample. 53 | Args: 54 | dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N] 55 | labels: pytorch LongTensor, with shape [N] 56 | return_inds: whether to return the indices. Save time if `False`(?) 57 | Returns: 58 | dist_ap: pytorch Variable, distance(anchor, positive); shape [N] 59 | dist_an: pytorch Variable, distance(anchor, negative); shape [N] 60 | p_inds: pytorch LongTensor, with shape [N]; 61 | indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1 62 | n_inds: pytorch LongTensor, with shape [N]; 63 | indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1 64 | NOTE: Only consider the case in which all labels have same num of samples, 65 | thus we can cope with all anchors in parallel. 66 | """ 67 | 68 | assert len(dist_mat.size()) == 2 69 | assert dist_mat.size(0) == dist_mat.size(1) 70 | N = dist_mat.size(0) 71 | 72 | # shape [N, N] 73 | is_pos = labels.expand(N, N).eq(labels.expand(N, N).t()) 74 | is_neg = labels.expand(N, N).ne(labels.expand(N, N).t()) 75 | 76 | # `dist_ap` means distance(anchor, positive) 77 | # both `dist_ap` and `relative_p_inds` with shape [N, 1] 78 | dist_ap, relative_p_inds = torch.max( 79 | dist_mat[is_pos].contiguous().view(N, -1), 1, keepdim=True) 80 | # print(dist_mat[is_pos].shape) 81 | # `dist_an` means distance(anchor, negative) 82 | # both `dist_an` and `relative_n_inds` with shape [N, 1] 83 | dist_an, relative_n_inds = torch.min( 84 | dist_mat[is_neg].contiguous().view(N, -1), 1, keepdim=True) 85 | # shape [N] 86 | dist_ap = dist_ap.squeeze(1) 87 | dist_an = dist_an.squeeze(1) 88 | 89 | if return_inds: 90 | # shape [N, N] 91 | ind = (labels.new().resize_as_(labels) 92 | .copy_(torch.arange(0, N).long()) 93 | .unsqueeze(0).expand(N, N)) 94 | # shape [N, 1] 95 | p_inds = torch.gather( 96 | ind[is_pos].contiguous().view(N, -1), 1, relative_p_inds.data) 97 | n_inds = torch.gather( 98 | ind[is_neg].contiguous().view(N, -1), 1, relative_n_inds.data) 99 | # shape [N] 100 | p_inds = p_inds.squeeze(1) 101 | n_inds = n_inds.squeeze(1) 102 | return dist_ap, dist_an, p_inds, n_inds 103 | 104 | return dist_ap, dist_an 105 | 106 | 107 | class TripletLoss(object): 108 | """ 109 | Triplet loss using HARDER example mining, 110 | modified based on original triplet loss using hard example mining 111 | """ 112 | 113 | def __init__(self, margin=None, hard_factor=0.0): 114 | self.margin = margin 115 | self.hard_factor = hard_factor 116 | if margin is not None: 117 | self.ranking_loss = nn.MarginRankingLoss(margin=margin) 118 | else: 119 | self.ranking_loss = nn.SoftMarginLoss() 120 | 121 | def __call__(self, global_feat, labels, normalize_feature=False): 122 | if normalize_feature: 123 | global_feat = normalize(global_feat, axis=-1) 124 | dist_mat = euclidean_dist(global_feat, global_feat) 125 | dist_ap, dist_an = hard_example_mining(dist_mat, labels) 126 | 127 | dist_ap *= (1.0 + self.hard_factor) 128 | dist_an *= (1.0 - self.hard_factor) 129 | 130 | y = dist_an.new().resize_as_(dist_an).fill_(1) 131 | if self.margin is not None: 132 | loss = self.ranking_loss(dist_an, dist_ap, y) 133 | else: 134 | loss = self.ranking_loss(dist_an - dist_ap, y) 135 | return loss, dist_ap, dist_an 136 | 137 | 138 | -------------------------------------------------------------------------------- /model/__init__.py: -------------------------------------------------------------------------------- 1 | from .make_model import make_model -------------------------------------------------------------------------------- /model/backbones/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhangguiwei610/PHA/fca6ece297036646c99d4af6496df237b94a8606/model/backbones/__init__.py -------------------------------------------------------------------------------- /model/backbones/resnet.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | from torch import nn 5 | 6 | 7 | def conv3x3(in_planes, out_planes, stride=1): 8 | """3x3 convolution with padding""" 9 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 10 | padding=1, bias=False) 11 | 12 | 13 | class BasicBlock(nn.Module): 14 | expansion = 1 15 | 16 | def __init__(self, inplanes, planes, stride=1, downsample=None): 17 | super(BasicBlock, self).__init__() 18 | self.conv1 = conv3x3(inplanes, planes, stride) 19 | self.bn1 = nn.BatchNorm2d(planes) 20 | self.relu = nn.ReLU(inplace=True) 21 | self.conv2 = conv3x3(planes, planes) 22 | self.bn2 = nn.BatchNorm2d(planes) 23 | self.downsample = downsample 24 | self.stride = stride 25 | 26 | def forward(self, x): 27 | residual = x 28 | 29 | out = self.conv1(x) 30 | out = self.bn1(out) 31 | out = self.relu(out) 32 | 33 | out = self.conv2(out) 34 | out = self.bn2(out) 35 | 36 | if self.downsample is not None: 37 | residual = self.downsample(x) 38 | 39 | out += residual 40 | out = self.relu(out) 41 | 42 | return out 43 | 44 | 45 | class Bottleneck(nn.Module): 46 | expansion = 4 47 | 48 | def __init__(self, inplanes, planes, stride=1, downsample=None): 49 | super(Bottleneck, self).__init__() 50 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 51 | self.bn1 = nn.BatchNorm2d(planes) 52 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 53 | padding=1, bias=False) 54 | self.bn2 = nn.BatchNorm2d(planes) 55 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 56 | self.bn3 = nn.BatchNorm2d(planes * 4) 57 | self.relu = nn.ReLU(inplace=True) 58 | self.downsample = downsample 59 | self.stride = stride 60 | 61 | def forward(self, x): 62 | residual = x 63 | 64 | out = self.conv1(x) 65 | out = self.bn1(out) 66 | out = self.relu(out) 67 | 68 | out = self.conv2(out) 69 | out = self.bn2(out) 70 | out = self.relu(out) 71 | 72 | out = self.conv3(out) 73 | out = self.bn3(out) 74 | 75 | if self.downsample is not None: 76 | residual = self.downsample(x) 77 | 78 | out += residual 79 | out = self.relu(out) 80 | 81 | return out 82 | 83 | 84 | class ResNet(nn.Module): 85 | def __init__(self, last_stride=2, block=Bottleneck,layers=[3, 4, 6, 3]): 86 | self.inplanes = 64 87 | super().__init__() 88 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 89 | bias=False) 90 | self.bn1 = nn.BatchNorm2d(64) 91 | # self.relu = nn.ReLU(inplace=True) # add missed relu 92 | self.maxpool = nn.MaxPool2d(kernel_size=2, stride=None, padding=0) 93 | self.layer1 = self._make_layer(block, 64, layers[0]) 94 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 95 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 96 | self.layer4 = self._make_layer(block, 512, layers[3], stride=last_stride) 97 | 98 | def _make_layer(self, block, planes, blocks, stride=1): 99 | downsample = None 100 | if stride != 1 or self.inplanes != planes * block.expansion: 101 | downsample = nn.Sequential( 102 | nn.Conv2d(self.inplanes, planes * block.expansion, 103 | kernel_size=1, stride=stride, bias=False), 104 | nn.BatchNorm2d(planes * block.expansion), 105 | ) 106 | 107 | layers = [] 108 | layers.append(block(self.inplanes, planes, stride, downsample)) 109 | self.inplanes = planes * block.expansion 110 | for i in range(1, blocks): 111 | layers.append(block(self.inplanes, planes)) 112 | 113 | return nn.Sequential(*layers) 114 | 115 | def forward(self, x, cam_label=None): 116 | x = self.conv1(x) 117 | x = self.bn1(x) 118 | # x = self.relu(x) # add missed relu 119 | x = self.maxpool(x) 120 | x = self.layer1(x) 121 | x = self.layer2(x) 122 | x = self.layer3(x) 123 | x = self.layer4(x) 124 | 125 | return x 126 | 127 | def load_param(self, model_path): 128 | param_dict = torch.load(model_path) 129 | for i in param_dict: 130 | if 'fc' in i: 131 | continue 132 | self.state_dict()[i].copy_(param_dict[i]) 133 | 134 | def random_init(self): 135 | for m in self.modules(): 136 | if isinstance(m, nn.Conv2d): 137 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 138 | m.weight.data.normal_(0, math.sqrt(2. / n)) 139 | elif isinstance(m, nn.BatchNorm2d): 140 | m.weight.data.fill_(1) 141 | m.bias.data.zero_() -------------------------------------------------------------------------------- /processor/__init__.py: -------------------------------------------------------------------------------- 1 | from .processor import do_train, do_inference -------------------------------------------------------------------------------- /processor/processor.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import time 4 | import torch 5 | import torch.nn as nn 6 | from utils.meter import AverageMeter 7 | from utils.metrics import R1_mAP_eval 8 | from torch.cuda import amp 9 | import torch.distributed as dist 10 | from PIL import Image 11 | from torchvision import transforms as trans 12 | import numpy 13 | def dequantize(image, q_table): 14 | """[summary] 15 | TODO: Add discription 16 | Args: 17 | image ([type]): [description] 18 | q_table ([type]): [description] 19 | 20 | Returns: 21 | [type]: [description] 22 | """ 23 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 24 | image = image.to(device) 25 | q_table = q_table.to(device) 26 | dequantitize_img = image * q_table 27 | return dequantitize_img 28 | 29 | def phi_diff(x, alpha): 30 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 31 | x = x.to(device) 32 | alpha = torch.where(alpha >= 2.0, torch.tensor([2.0]).cuda(), alpha) 33 | s = 1/(1-alpha).to(device) 34 | k = torch.log(2/alpha -1).to(device) 35 | phi_x = torch.tanh((x - (torch.floor(x) + 0.5)) * k) * s 36 | x_ = (phi_x + 1)/2 + torch.floor(x) 37 | return x_ 38 | def quantize(image, q_table,alpha): 39 | """[summary] 40 | TODO: add disciption. 41 | 42 | Args: 43 | image ([type]): [description] 44 | q_table ([type]): [description] 45 | """ 46 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 47 | image = image.to(device) 48 | q_table = q_table.to(device) 49 | pre_img = image/(q_table) 50 | after_img = phi_diff(pre_img, alpha) 51 | # after_img = sgn(after_img) 52 | # after_img = torch.round(pre_img) + torch.empty_like(pre_img).uniform_(0.0, 1.0) 53 | # diff = after_img - pre_img 54 | # print("Max difference: ", torch.max(diff)) 55 | # image = torch.round(image) 56 | # image = diff_round(image) 57 | # after_img = diff_round(pre_img) 58 | return after_img 59 | 60 | from timm.data.random_erasing import RandomErasing 61 | 62 | 63 | def do_train(cfg, 64 | model, 65 | center_criterion, 66 | train_loader, 67 | val_loader, 68 | optimizer, 69 | optimizer_center, 70 | scheduler, 71 | loss_fn, 72 | num_query, local_rank): 73 | log_period = cfg.SOLVER.LOG_PERIOD 74 | checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD 75 | eval_period = cfg.SOLVER.EVAL_PERIOD 76 | 77 | device = "cuda" 78 | epochs = cfg.SOLVER.MAX_EPOCHS 79 | 80 | logger = logging.getLogger("transreid.train") 81 | logger.info('start training') 82 | _LOCAL_PROCESS_GROUP = None 83 | if device: 84 | model.to(local_rank) 85 | if torch.cuda.device_count() > 1 and cfg.MODEL.DIST_TRAIN: 86 | print('Using {} GPUs for training'.format(torch.cuda.device_count())) 87 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True) 88 | 89 | loss_meter = AverageMeter() 90 | acc_meter = AverageMeter() 91 | 92 | evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM) 93 | scaler = amp.GradScaler() 94 | # train 95 | from pytorch_wavelets import DWTForward, DWTInverse # (or import DWT, IDWT) 96 | # J为分解的层次数,wave表示使用的变换方法 97 | xfm = DWTForward(J=1, mode='zero', wave='haar') # Accepts all wave types available to PyWavelets 98 | ifm = DWTInverse(mode='zero', wave='haar') 99 | transform = trans.Compose([ 100 | trans.Resize([256, 128], interpolation=3), 101 | trans.ToTensor() 102 | ]) 103 | 104 | 105 | transform_norm = trans.Compose([ 106 | trans.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD), 107 | RandomErasing(probability=cfg.INPUT.RE_PROB, mode='pixel', max_count=1, device='cpu'), 108 | # RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN) 109 | ]) 110 | 111 | for epoch in range(1, epochs + 1): 112 | start_time = time.time() 113 | loss_meter.reset() 114 | acc_meter.reset() 115 | evaluator.reset() 116 | scheduler.step(epoch) 117 | model.train() 118 | for n_iter, (img, vid, target_cam, target_view,img_paths) in enumerate(train_loader): 119 | factor = min(epoch * 0.05, 0.5) 120 | topk_num=int(128*64*0.4) 121 | wavelet_list = [] 122 | for per_path in img_paths: 123 | img_tmp = Image.open(per_path) 124 | img_tmp = transform(img_tmp) 125 | wavelet_list.append(img_tmp) 126 | wavelet_list = torch.stack(wavelet_list) # bs,3,256,128 127 | x_dwt_l, x_dwt_h = xfm(wavelet_list) # input bs,3,256,128 Yl bs,3,128,64 Yh[0] bs,c,3,128,64 128 | high_frequency = x_dwt_h[0].sum(dim=2) # x_dwt_h[0] bs,c,3,128,64，output bs,c,128,64 129 | high_frequency = torch.norm(high_frequency, dim=1) # bs,128,64 130 | 131 | cluster = high_frequency.flatten(1) # bs,128*64 132 | _, topk = cluster.topk(topk_num, dim=1, largest=True, sorted=True) # bs,256 133 | mask_map = torch.zeros_like(cluster) 134 | idx_batch = torch.arange(len(img_paths))[:, None].expand(-1, topk_num) 135 | mask_map[idx_batch, topk] = 255 136 | mask_map = mask_map.reshape(len(img_paths), 128, 64) # bs,128,64 137 | mask_map_adv=mask_map 138 | mask_map = torch.nn.functional.interpolate(mask_map[None, :, :, :], size=(21, 10)).squeeze(0) # bs,21,10 139 | 140 | # adv 141 | # Value for quantization range 142 | alpha_range = [0.1, 1e-20] 143 | alpha = torch.tensor(alpha_range[0]).cuda() 144 | 145 | q_ini_table = numpy.empty((32, 128, 64), dtype=numpy.float32) 146 | q_ini_table.fill(5) 147 | q_ini_table=numpy.where(mask_map_adv>0,q_ini_table,numpy.ones_like(q_ini_table)) 148 | q_tables = {"y": torch.from_numpy(q_ini_table).cuda(), 149 | "cb": torch.from_numpy(q_ini_table).cuda(), 150 | "cr": torch.from_numpy(q_ini_table).cuda()} 151 | alpha = alpha.to('cuda') 152 | 153 | 154 | x_dwt_l_q = 255.0 * x_dwt_l.clone().detach().to('cuda') # bs,c,h,w 155 | x_dwt_l_q = x_dwt_l_q.permute(0, 2, 3, 1) # bs,h,w,c 156 | components = {'y': x_dwt_l_q[:, :, :, 0], 'cb': x_dwt_l_q[:, :, :, 1], 157 | 'cr': x_dwt_l_q[:, :, :, 2]} # y,cb,cr bs,h,w 158 | # q_tables["y"].requires_grad = True 159 | # q_tables["cb"].requires_grad = True 160 | # q_tables["cr"].requires_grad = True 161 | upresults = {} 162 | for k in components.keys(): 163 | comp = quantize(components[k], q_tables[k], alpha) # output bs，784，8，8 164 | comp = dequantize(comp, q_tables[k]) # output bs，784，8，8 165 | upresults[k] = comp 166 | rgb_images = torch.cat( 167 | [upresults['y'].unsqueeze(3), upresults['cb'].unsqueeze(3), upresults['cr'].unsqueeze(3)], dim=3) 168 | rgb_images = rgb_images.permute(0, 3, 1, 2) / 255. # bs,3,128,64 169 | 170 | img_adv = ifm((rgb_images.cpu(), x_dwt_h)).cuda() 171 | img_adv = transform_norm(img_adv) 172 | 173 | 174 | 175 | 176 | optimizer.zero_grad() 177 | optimizer_center.zero_grad() 178 | img = img.to(device) 179 | target = vid.to(device) 180 | target_cam = target_cam.to(device) 181 | target_view = target_view.to(device) 182 | with amp.autocast(enabled=True): 183 | score, feat ,mask,mask_adv,loss_contra= model([img,img_adv], target, cam_label=target_cam, view_label=target_view ,mask=mask_map) 184 | 185 | loss = loss_fn(score, feat, target, target_cam) 186 | loss_mask=loss_fn(mask[0],mask[1],target,target_cam)*factor 187 | 188 | loss_mask_adv = loss_fn(mask_adv[0], mask_adv[1], target, target_cam) * factor 189 | 190 | 191 | scaler.scale(loss+loss_mask+loss_mask_adv+loss_contra).backward() 192 | 193 | 194 | scaler.step(optimizer) 195 | scaler.update() 196 | 197 | if 'center' in cfg.MODEL.METRIC_LOSS_TYPE: 198 | for param in center_criterion.parameters(): 199 | param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT) 200 | scaler.step(optimizer_center) 201 | scaler.update() 202 | if isinstance(score, list): 203 | acc = (score[0].max(1)[1] == target).float().mean() 204 | else: 205 | acc = (score.max(1)[1] == target).float().mean() 206 | 207 | loss_meter.update(loss.item(), img.shape[0]) 208 | acc_meter.update(acc, 1) 209 | 210 | torch.cuda.synchronize() 211 | if (n_iter + 1) % log_period == 0: 212 | logger.info("Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Loss_mask :{:.3f}, Acc: {:.3f}, Base Lr: {:.2e}" 213 | .format(epoch, (n_iter + 1), len(train_loader), 214 | loss_meter.avg,loss_mask.item(), acc_meter.avg, scheduler._get_lr(epoch)[0])) 215 | 216 | end_time = time.time() 217 | time_per_batch = (end_time - start_time) / (n_iter + 1) 218 | if cfg.MODEL.DIST_TRAIN: 219 | pass 220 | else: 221 | logger.info("Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]" 222 | .format(epoch, time_per_batch, train_loader.batch_size / time_per_batch)) 223 | 224 | if epoch % checkpoint_period == 0: 225 | if cfg.MODEL.DIST_TRAIN: 226 | if dist.get_rank() == 0: 227 | torch.save(model.state_dict(), 228 | os.path.join(cfg.OUTPUT_DIR, cfg.MODEL.NAME + '_{}.pth'.format(epoch))) 229 | else: 230 | torch.save(model.state_dict(), 231 | os.path.join(cfg.OUTPUT_DIR, cfg.MODEL.NAME + '_{}.pth'.format(epoch))) 232 | 233 | if epoch % 40 == 0: 234 | if cfg.MODEL.DIST_TRAIN: 235 | if dist.get_rank() == 0: 236 | model.eval() 237 | for n_iter, (img, vid, camid, camids, target_view, _) in enumerate(val_loader): 238 | with torch.no_grad(): 239 | img = img.to(device) 240 | camids = camids.to(device) 241 | target_view = target_view.to(device) 242 | feat = model(img, cam_label=camids, view_label=target_view) 243 | evaluator.update((feat, vid, camid)) 244 | cmc, mAP, _, _, _, _, _ = evaluator.compute() 245 | logger.info("Validation Results - Epoch: {}".format(epoch)) 246 | logger.info("mAP: {:.1%}".format(mAP)) 247 | for r in [1, 5, 10]: 248 | logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1])) 249 | torch.cuda.empty_cache() 250 | else: 251 | model.eval() 252 | for n_iter, (img, vid, camid, camids, target_view, _) in enumerate(val_loader): 253 | with torch.no_grad(): 254 | img = img.to(device) 255 | camids = camids.to(device) 256 | target_view = target_view.to(device) 257 | feat = model(img, cam_label=camids, view_label=target_view) 258 | evaluator.update((feat, vid, camid)) 259 | cmc, mAP, _, _, _, _, _ = evaluator.compute() 260 | logger.info("Validation Results - Epoch: {}".format(epoch)) 261 | logger.info("mAP: {:.1%}".format(mAP)) 262 | for r in [1, 5, 10]: 263 | logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1])) 264 | torch.cuda.empty_cache() 265 | 266 | 267 | def do_inference(cfg, 268 | model, 269 | val_loader, 270 | num_query): 271 | device = "cuda" 272 | logger = logging.getLogger("transreid.test") 273 | logger.info("Enter inferencing") 274 | 275 | evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM) 276 | 277 | evaluator.reset() 278 | 279 | if device: 280 | if torch.cuda.device_count() > 1: 281 | print('Using {} GPUs for inference'.format(torch.cuda.device_count())) 282 | model = nn.DataParallel(model) 283 | model.to(device) 284 | 285 | model.eval() 286 | img_path_list = [] 287 | 288 | for n_iter, (img, pid, camid, camids, target_view, imgpath) in enumerate(val_loader): 289 | with torch.no_grad(): 290 | img = img.to(device) 291 | camids = camids.to(device) 292 | target_view = target_view.to(device) 293 | feat = model(img, cam_label=camids, view_label=target_view) 294 | evaluator.update((feat, pid, camid)) 295 | img_path_list.extend(imgpath) 296 | 297 | cmc, mAP, _, _, _, _, _ = evaluator.compute() 298 | logger.info("Validation Results ") 299 | logger.info("mAP: {:.1%}".format(mAP)) 300 | for r in [1, 5, 10]: 301 | logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1])) 302 | return cmc[0], cmc[4] 303 | 304 | 305 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch 2 | torchvision 3 | timm 4 | yacs 5 | opencv-python -------------------------------------------------------------------------------- /solver/__init__.py: -------------------------------------------------------------------------------- 1 | from .lr_scheduler import WarmupMultiStepLR 2 | from .make_optimizer import make_optimizer -------------------------------------------------------------------------------- /solver/cosine_lr.py: -------------------------------------------------------------------------------- 1 | """ Cosine Scheduler 2 | 3 | Cosine LR schedule with warmup, cycle/restarts, noise. 4 | 5 | Hacked together by / Copyright 2020 Ross Wightman 6 | """ 7 | import logging 8 | import math 9 | import torch 10 | 11 | from .scheduler import Scheduler 12 | 13 | 14 | _logger = logging.getLogger(__name__) 15 | 16 | 17 | class CosineLRScheduler(Scheduler): 18 | """ 19 | Cosine decay with restarts. 20 | This is described in the paper https://arxiv.org/abs/1608.03983. 21 | 22 | Inspiration from 23 | https://github.com/allenai/allennlp/blob/master/allennlp/training/learning_rate_schedulers/cosine.py 24 | """ 25 | 26 | def __init__(self, 27 | optimizer: torch.optim.Optimizer, 28 | t_initial: int, 29 | t_mul: float = 1., 30 | lr_min: float = 0., 31 | decay_rate: float = 1., 32 | warmup_t=0, 33 | warmup_lr_init=0, 34 | warmup_prefix=False, 35 | cycle_limit=0, 36 | t_in_epochs=True, 37 | noise_range_t=None, 38 | noise_pct=0.67, 39 | noise_std=1.0, 40 | noise_seed=42, 41 | initialize=True) -> None: 42 | super().__init__( 43 | optimizer, param_group_field="lr", 44 | noise_range_t=noise_range_t, noise_pct=noise_pct, noise_std=noise_std, noise_seed=noise_seed, 45 | initialize=initialize) 46 | 47 | assert t_initial > 0 48 | assert lr_min >= 0 49 | if t_initial == 1 and t_mul == 1 and decay_rate == 1: 50 | _logger.warning("Cosine annealing scheduler will have no effect on the learning " 51 | "rate since t_initial = t_mul = eta_mul = 1.") 52 | self.t_initial = t_initial 53 | self.t_mul = t_mul 54 | self.lr_min = lr_min 55 | self.decay_rate = decay_rate 56 | self.cycle_limit = cycle_limit 57 | self.warmup_t = warmup_t 58 | self.warmup_lr_init = warmup_lr_init 59 | self.warmup_prefix = warmup_prefix 60 | self.t_in_epochs = t_in_epochs 61 | if self.warmup_t: 62 | self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values] 63 | super().update_groups(self.warmup_lr_init) 64 | else: 65 | self.warmup_steps = [1 for _ in self.base_values] 66 | 67 | def _get_lr(self, t): 68 | if t < self.warmup_t: 69 | lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps] 70 | else: 71 | if self.warmup_prefix: 72 | t = t - self.warmup_t 73 | 74 | if self.t_mul != 1: 75 | i = math.floor(math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul)) 76 | t_i = self.t_mul ** i * self.t_initial 77 | t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial 78 | else: 79 | i = t // self.t_initial 80 | t_i = self.t_initial 81 | t_curr = t - (self.t_initial * i) 82 | 83 | gamma = self.decay_rate ** i 84 | lr_min = self.lr_min * gamma 85 | lr_max_values = [v * gamma for v in self.base_values] 86 | 87 | if self.cycle_limit == 0 or (self.cycle_limit > 0 and i < self.cycle_limit): 88 | lrs = [ 89 | lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t_curr / t_i)) for lr_max in lr_max_values 90 | ] 91 | else: 92 | lrs = [self.lr_min for _ in self.base_values] 93 | 94 | return lrs 95 | 96 | def get_epoch_values(self, epoch: int): 97 | if self.t_in_epochs: 98 | return self._get_lr(epoch) 99 | else: 100 | return None 101 | 102 | def get_update_values(self, num_updates: int): 103 | if not self.t_in_epochs: 104 | return self._get_lr(num_updates) 105 | else: 106 | return None 107 | 108 | def get_cycle_length(self, cycles=0): 109 | if not cycles: 110 | cycles = self.cycle_limit 111 | cycles = max(1, cycles) 112 | if self.t_mul == 1.0: 113 | return self.t_initial * cycles 114 | else: 115 | return int(math.floor(-self.t_initial * (self.t_mul ** cycles - 1) / (1 - self.t_mul))) 116 | -------------------------------------------------------------------------------- /solver/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: liaoxingyu 4 | @contact: sherlockliao01@gmail.com 5 | """ 6 | from bisect import bisect_right 7 | import torch 8 | 9 | 10 | # FIXME ideally this would be achieved with a CombinedLRScheduler, 11 | # separating MultiStepLR with WarmupLR 12 | # but the current LRScheduler design doesn't allow it 13 | 14 | class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler): 15 | def __init__( 16 | self, 17 | optimizer, 18 | milestones, # steps 19 | gamma=0.1, 20 | warmup_factor=1.0 / 3, 21 | warmup_iters=500, 22 | warmup_method="linear", 23 | last_epoch=-1, 24 | ): 25 | if not list(milestones) == sorted(milestones): 26 | raise ValueError( 27 | "Milestones should be a list of" " increasing integers. Got {}", 28 | milestones, 29 | ) 30 | 31 | if warmup_method not in ("constant", "linear"): 32 | raise ValueError( 33 | "Only 'constant' or 'linear' warmup_method accepted" 34 | "got {}".format(warmup_method) 35 | ) 36 | self.milestones = milestones 37 | self.gamma = gamma 38 | self.warmup_factor = warmup_factor 39 | self.warmup_iters = warmup_iters 40 | self.warmup_method = warmup_method 41 | super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch) 42 | 43 | def _get_lr(self): 44 | warmup_factor = 1 45 | if self.last_epoch < self.warmup_iters: 46 | if self.warmup_method == "constant": 47 | warmup_factor = self.warmup_factor 48 | elif self.warmup_method == "linear": 49 | alpha = self.last_epoch / self.warmup_iters 50 | warmup_factor = self.warmup_factor * (1 - alpha) + alpha 51 | return [ 52 | base_lr 53 | * warmup_factor 54 | * self.gamma ** bisect_right(self.milestones, self.last_epoch) 55 | for base_lr in self.base_lrs 56 | ] 57 | -------------------------------------------------------------------------------- /solver/make_optimizer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def make_optimizer(cfg, model, center_criterion): 5 | params = [] 6 | for key, value in model.named_parameters(): 7 | if not value.requires_grad: 8 | continue 9 | lr = cfg.SOLVER.BASE_LR 10 | weight_decay = cfg.SOLVER.WEIGHT_DECAY 11 | if "bias" in key: 12 | lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR 13 | weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS 14 | if cfg.SOLVER.LARGE_FC_LR: 15 | if "classifier" in key or "arcface" in key: 16 | lr = cfg.SOLVER.BASE_LR * 2 17 | print('Using two times learning rate for fc ') 18 | 19 | params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}] 20 | 21 | if cfg.SOLVER.OPTIMIZER_NAME == 'SGD': 22 | optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params, momentum=cfg.SOLVER.MOMENTUM) 23 | elif cfg.SOLVER.OPTIMIZER_NAME == 'AdamW': 24 | optimizer = torch.optim.AdamW(params, lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY) 25 | else: 26 | optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params) 27 | optimizer_center = torch.optim.SGD(center_criterion.parameters(), lr=cfg.SOLVER.CENTER_LR) 28 | 29 | return optimizer, optimizer_center 30 | -------------------------------------------------------------------------------- /solver/scheduler.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Any 2 | 3 | import torch 4 | 5 | 6 | class Scheduler: 7 | """ Parameter Scheduler Base Class 8 | A scheduler base class that can be used to schedule any optimizer parameter groups. 9 | 10 | Unlike the builtin PyTorch schedulers, this is intended to be consistently called 11 | * At the END of each epoch, before incrementing the epoch count, to calculate next epoch's value 12 | * At the END of each optimizer update, after incrementing the update count, to calculate next update's value 13 | 14 | The schedulers built on this should try to remain as stateless as possible (for simplicity). 15 | 16 | This family of schedulers is attempting to avoid the confusion of the meaning of 'last_epoch' 17 | and -1 values for special behaviour. All epoch and update counts must be tracked in the training 18 | code and explicitly passed in to the schedulers on the corresponding step or step_update call. 19 | 20 | Based on ideas from: 21 | * https://github.com/pytorch/fairseq/tree/master/fairseq/optim/lr_scheduler 22 | * https://github.com/allenai/allennlp/tree/master/allennlp/training/learning_rate_schedulers 23 | """ 24 | 25 | def __init__(self, 26 | optimizer: torch.optim.Optimizer, 27 | param_group_field: str, 28 | noise_range_t=None, 29 | noise_type='normal', 30 | noise_pct=0.67, 31 | noise_std=1.0, 32 | noise_seed=None, 33 | initialize: bool = True) -> None: 34 | self.optimizer = optimizer 35 | self.param_group_field = param_group_field 36 | self._initial_param_group_field = f"initial_{param_group_field}" 37 | if initialize: 38 | for i, group in enumerate(self.optimizer.param_groups): 39 | if param_group_field not in group: 40 | raise KeyError(f"{param_group_field} missing from param_groups[{i}]") 41 | group.setdefault(self._initial_param_group_field, group[param_group_field]) 42 | else: 43 | for i, group in enumerate(self.optimizer.param_groups): 44 | if self._initial_param_group_field not in group: 45 | raise KeyError(f"{self._initial_param_group_field} missing from param_groups[{i}]") 46 | self.base_values = [group[self._initial_param_group_field] for group in self.optimizer.param_groups] 47 | self.metric = None # any point to having this for all? 48 | self.noise_range_t = noise_range_t 49 | self.noise_pct = noise_pct 50 | self.noise_type = noise_type 51 | self.noise_std = noise_std 52 | self.noise_seed = noise_seed if noise_seed is not None else 42 53 | self.update_groups(self.base_values) 54 | 55 | def state_dict(self) -> Dict[str, Any]: 56 | return {key: value for key, value in self.__dict__.items() if key != 'optimizer'} 57 | 58 | def load_state_dict(self, state_dict: Dict[str, Any]) -> None: 59 | self.__dict__.update(state_dict) 60 | 61 | def get_epoch_values(self, epoch: int): 62 | return None 63 | 64 | def get_update_values(self, num_updates: int): 65 | return None 66 | 67 | def step(self, epoch: int, metric: float = None) -> None: 68 | self.metric = metric 69 | values = self.get_epoch_values(epoch) 70 | if values is not None: 71 | values = self._add_noise(values, epoch) 72 | self.update_groups(values) 73 | 74 | def step_update(self, num_updates: int, metric: float = None): 75 | self.metric = metric 76 | values = self.get_update_values(num_updates) 77 | if values is not None: 78 | values = self._add_noise(values, num_updates) 79 | self.update_groups(values) 80 | 81 | def update_groups(self, values): 82 | if not isinstance(values, (list, tuple)): 83 | values = [values] * len(self.optimizer.param_groups) 84 | for param_group, value in zip(self.optimizer.param_groups, values): 85 | param_group[self.param_group_field] = value 86 | 87 | def _add_noise(self, lrs, t): 88 | if self.noise_range_t is not None: 89 | if isinstance(self.noise_range_t, (list, tuple)): 90 | apply_noise = self.noise_range_t[0] <= t < self.noise_range_t[1] 91 | else: 92 | apply_noise = t >= self.noise_range_t 93 | if apply_noise: 94 | g = torch.Generator() 95 | g.manual_seed(self.noise_seed + t) 96 | if self.noise_type == 'normal': 97 | while True: 98 | # resample if noise out of percent limit, brute force but shouldn't spin much 99 | noise = torch.randn(1, generator=g).item() 100 | if abs(noise) < self.noise_pct: 101 | break 102 | else: 103 | noise = 2 * (torch.rand(1, generator=g).item() - 0.5) * self.noise_pct 104 | lrs = [v + v * noise for v in lrs] 105 | return lrs 106 | -------------------------------------------------------------------------------- /solver/scheduler_factory.py: -------------------------------------------------------------------------------- 1 | """ Scheduler Factory 2 | Hacked together by / Copyright 2020 Ross Wightman 3 | """ 4 | from .cosine_lr import CosineLRScheduler 5 | 6 | 7 | def create_scheduler(cfg, optimizer): 8 | num_epochs = cfg.SOLVER.MAX_EPOCHS 9 | # type 1 10 | # lr_min = 0.01 * cfg.SOLVER.BASE_LR 11 | # warmup_lr_init = 0.001 * cfg.SOLVER.BASE_LR 12 | # type 2 13 | lr_min = 0.002 * cfg.SOLVER.BASE_LR 14 | warmup_lr_init = 0.01 * cfg.SOLVER.BASE_LR 15 | # type 3 16 | # lr_min = 0.001 * cfg.SOLVER.BASE_LR 17 | # warmup_lr_init = 0.01 * cfg.SOLVER.BASE_LR 18 | 19 | warmup_t = cfg.SOLVER.WARMUP_EPOCHS 20 | noise_range = None 21 | 22 | lr_scheduler = CosineLRScheduler( 23 | optimizer, 24 | t_initial=num_epochs, 25 | lr_min=lr_min, 26 | t_mul= 1., 27 | decay_rate=0.1, 28 | warmup_lr_init=warmup_lr_init, 29 | warmup_t=warmup_t, 30 | cycle_limit=1, 31 | t_in_epochs=True, 32 | noise_range_t=noise_range, 33 | noise_pct= 0.67, 34 | noise_std= 1., 35 | noise_seed=42, 36 | ) 37 | 38 | return lr_scheduler 39 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import os 2 | from config import cfg 3 | import argparse 4 | from datasets import make_dataloader 5 | from model import make_model 6 | from processor import do_inference 7 | from utils.logger import setup_logger 8 | 9 | 10 | if __name__ == "__main__": 11 | parser = argparse.ArgumentParser(description="ReID Baseline Training") 12 | parser.add_argument( 13 | "--config_file", default="", help="path to config file", type=str 14 | ) 15 | parser.add_argument("opts", help="Modify config options using the command-line", default=None, 16 | nargs=argparse.REMAINDER) 17 | 18 | args = parser.parse_args() 19 | 20 | 21 | 22 | if args.config_file != "": 23 | cfg.merge_from_file(args.config_file) 24 | cfg.merge_from_list(args.opts) 25 | cfg.freeze() 26 | 27 | output_dir = cfg.OUTPUT_DIR 28 | if output_dir and not os.path.exists(output_dir): 29 | os.makedirs(output_dir) 30 | 31 | logger = setup_logger("transreid", output_dir, if_train=False) 32 | logger.info(args) 33 | 34 | if args.config_file != "": 35 | logger.info("Loaded configuration file {}".format(args.config_file)) 36 | with open(args.config_file, 'r') as cf: 37 | config_str = "\n" + cf.read() 38 | logger.info(config_str) 39 | logger.info("Running with config:\n{}".format(cfg)) 40 | 41 | os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID 42 | 43 | train_loader, train_loader_normal, val_loader, num_query, num_classes, camera_num, view_num = make_dataloader(cfg) 44 | 45 | model = make_model(cfg, num_class=num_classes, camera_num=camera_num, view_num = view_num) 46 | model.load_param(cfg.TEST.WEIGHT) 47 | 48 | if cfg.DATASETS.NAMES == 'VehicleID': 49 | for trial in range(10): 50 | train_loader, train_loader_normal, val_loader, num_query, num_classes, camera_num, view_num = make_dataloader(cfg) 51 | rank_1, rank5 = do_inference(cfg, 52 | model, 53 | val_loader, 54 | num_query) 55 | if trial == 0: 56 | all_rank_1 = rank_1 57 | all_rank_5 = rank5 58 | else: 59 | all_rank_1 = all_rank_1 + rank_1 60 | all_rank_5 = all_rank_5 + rank5 61 | 62 | logger.info("rank_1:{}, rank_5 {} : trial : {}".format(rank_1, rank5, trial)) 63 | logger.info("sum_rank_1:{:.1%}, sum_rank_5 {:.1%}".format(all_rank_1.sum()/10.0, all_rank_5.sum()/10.0)) 64 | else: 65 | do_inference(cfg, 66 | model, 67 | val_loader, 68 | num_query) 69 | 70 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | from utils.logger import setup_logger 2 | from datasets import make_dataloader 3 | from model import make_model 4 | from solver import make_optimizer 5 | from solver.scheduler_factory import create_scheduler 6 | from loss import make_loss 7 | from processor import do_train 8 | import random 9 | import torch 10 | import numpy as np 11 | import os 12 | import argparse 13 | # from timm.scheduler import create_scheduler 14 | from config import cfg 15 | 16 | def set_seed(seed): 17 | torch.manual_seed(seed) 18 | torch.cuda.manual_seed(seed) 19 | torch.cuda.manual_seed_all(seed) 20 | np.random.seed(seed) 21 | random.seed(seed) 22 | torch.backends.cudnn.deterministic = True 23 | torch.backends.cudnn.benchmark = True 24 | 25 | if __name__ == '__main__': 26 | 27 | parser = argparse.ArgumentParser(description="[CVPR2023] PHA Training") 28 | parser.add_argument( 29 | "--config_file", default="./configs/Cuhk03_labeled/vit_transreid_stride.yml", help="path to config file", type=str 30 | ) 31 | 32 | parser.add_argument("opts", help="Modify config options using the command-line", default=None, 33 | nargs=argparse.REMAINDER) 34 | parser.add_argument("--local_rank", default=0, type=int) 35 | args = parser.parse_args() 36 | 37 | if args.config_file != "": 38 | cfg.merge_from_file(args.config_file) 39 | cfg.merge_from_list(args.opts) 40 | cfg.freeze() 41 | 42 | set_seed(cfg.SOLVER.SEED) 43 | 44 | if cfg.MODEL.DIST_TRAIN: 45 | torch.cuda.set_device(args.local_rank) 46 | 47 | output_dir = cfg.OUTPUT_DIR 48 | if output_dir and not os.path.exists(output_dir): 49 | os.makedirs(output_dir) 50 | 51 | logger = setup_logger("transreid", output_dir, if_train=True) 52 | logger.info("Saving model in the path :{}".format(cfg.OUTPUT_DIR)) 53 | logger.info(args) 54 | 55 | if args.config_file != "": 56 | logger.info("Loaded configuration file {}".format(args.config_file)) 57 | with open(args.config_file, 'r') as cf: 58 | config_str = "\n" + cf.read() 59 | logger.info(config_str) 60 | logger.info("Running with config:\n{}".format(cfg)) 61 | 62 | if cfg.MODEL.DIST_TRAIN: 63 | torch.distributed.init_process_group(backend='nccl', init_method='env://') 64 | 65 | os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID 66 | train_loader, train_loader_normal, val_loader, num_query, num_classes, camera_num, view_num = make_dataloader(cfg) 67 | 68 | model = make_model(cfg, num_class=num_classes, camera_num=camera_num, view_num = view_num) 69 | 70 | loss_func, center_criterion = make_loss(cfg, num_classes=num_classes) 71 | 72 | optimizer, optimizer_center = make_optimizer(cfg, model, center_criterion) 73 | 74 | scheduler = create_scheduler(cfg, optimizer) 75 | 76 | do_train( 77 | cfg, 78 | model, 79 | center_criterion, 80 | train_loader, 81 | val_loader, 82 | optimizer, 83 | optimizer_center, 84 | scheduler, 85 | loss_func, 86 | num_query, args.local_rank 87 | ) 88 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zhangguiwei610/PHA/fca6ece297036646c99d4af6496df237b94a8606/utils/__init__.py -------------------------------------------------------------------------------- /utils/iotools.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: sherlock 4 | @contact: sherlockliao01@gmail.com 5 | """ 6 | 7 | import errno 8 | import json 9 | import os 10 | 11 | import os.path as osp 12 | 13 | 14 | def mkdir_if_missing(directory): 15 | if not osp.exists(directory): 16 | try: 17 | os.makedirs(directory) 18 | except OSError as e: 19 | if e.errno != errno.EEXIST: 20 | raise 21 | 22 | 23 | def check_isfile(path): 24 | isfile = osp.isfile(path) 25 | if not isfile: 26 | print("=> Warning: no file found at '{}' (ignored)".format(path)) 27 | return isfile 28 | 29 | 30 | def read_json(fpath): 31 | with open(fpath, 'r') as f: 32 | obj = json.load(f) 33 | return obj 34 | 35 | 36 | def write_json(obj, fpath): 37 | mkdir_if_missing(osp.dirname(fpath)) 38 | with open(fpath, 'w') as f: 39 | json.dump(obj, f, indent=4, separators=(',', ': ')) 40 | -------------------------------------------------------------------------------- /utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import sys 4 | import os.path as osp 5 | def setup_logger(name, save_dir, if_train): 6 | logger = logging.getLogger(name) 7 | logger.setLevel(logging.DEBUG) 8 | 9 | ch = logging.StreamHandler(stream=sys.stdout) 10 | ch.setLevel(logging.DEBUG) 11 | formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s") 12 | ch.setFormatter(formatter) 13 | logger.addHandler(ch) 14 | 15 | if save_dir: 16 | if not osp.exists(save_dir): 17 | os.makedirs(save_dir) 18 | if if_train: 19 | fh = logging.FileHandler(os.path.join(save_dir, "train_log.txt"), mode='w') 20 | else: 21 | fh = logging.FileHandler(os.path.join(save_dir, "test_log.txt"), mode='w') 22 | fh.setLevel(logging.DEBUG) 23 | fh.setFormatter(formatter) 24 | logger.addHandler(fh) 25 | 26 | return logger -------------------------------------------------------------------------------- /utils/meter.py: -------------------------------------------------------------------------------- 1 | class AverageMeter(object): 2 | """Computes and stores the average and current value""" 3 | 4 | def __init__(self): 5 | self.val = 0 6 | self.avg = 0 7 | self.sum = 0 8 | self.count = 0 9 | 10 | def reset(self): 11 | self.val = 0 12 | self.avg = 0 13 | self.sum = 0 14 | self.count = 0 15 | 16 | def update(self, val, n=1): 17 | self.val = val 18 | self.sum += val * n 19 | self.count += n 20 | self.avg = self.sum / self.count -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import os 4 | from utils.reranking import re_ranking 5 | 6 | 7 | def euclidean_distance(qf, gf): 8 | m = qf.shape[0] 9 | n = gf.shape[0] 10 | dist_mat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \ 11 | torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t() 12 | dist_mat.addmm_(1, -2, qf, gf.t()) 13 | return dist_mat.cpu().numpy() 14 | 15 | def cosine_similarity(qf, gf): 16 | epsilon = 0.00001 17 | dist_mat = qf.mm(gf.t()) 18 | qf_norm = torch.norm(qf, p=2, dim=1, keepdim=True) # mx1 19 | gf_norm = torch.norm(gf, p=2, dim=1, keepdim=True) # nx1 20 | qg_normdot = qf_norm.mm(gf_norm.t()) 21 | 22 | dist_mat = dist_mat.mul(1 / qg_normdot).cpu().numpy() 23 | dist_mat = np.clip(dist_mat, -1 + epsilon, 1 - epsilon) 24 | dist_mat = np.arccos(dist_mat) 25 | return dist_mat 26 | 27 | 28 | def eval_func(distmat, q_pids, g_pids, q_camids, g_camids, max_rank=50): 29 | """Evaluation with market1501 metric 30 | Key: for each query identity, its gallery images from the same camera view are discarded. 31 | """ 32 | num_q, num_g = distmat.shape 33 | # distmat g 34 | # q 1 3 2 4 35 | # 4 1 2 3 36 | if num_g < max_rank: 37 | max_rank = num_g 38 | print("Note: number of gallery samples is quite small, got {}".format(num_g)) 39 | indices = np.argsort(distmat, axis=1) 40 | # 0 2 1 3 41 | # 1 2 3 0 42 | matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32) 43 | # compute cmc curve for each query 44 | all_cmc = [] 45 | all_AP = [] 46 | num_valid_q = 0. # number of valid query 47 | for q_idx in range(num_q): 48 | # get query pid and camid 49 | q_pid = q_pids[q_idx] 50 | q_camid = q_camids[q_idx] 51 | 52 | # remove gallery samples that have the same pid and camid with query 53 | order = indices[q_idx] # select one row 54 | remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid) 55 | keep = np.invert(remove) 56 | 57 | # compute cmc curve 58 | # binary vector, positions with value 1 are correct matches 59 | orig_cmc = matches[q_idx][keep] 60 | if not np.any(orig_cmc): 61 | # this condition is true when query identity does not appear in gallery 62 | continue 63 | 64 | cmc = orig_cmc.cumsum() 65 | cmc[cmc > 1] = 1 66 | 67 | all_cmc.append(cmc[:max_rank]) 68 | num_valid_q += 1. 69 | 70 | # compute average precision 71 | # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision 72 | num_rel = orig_cmc.sum() 73 | tmp_cmc = orig_cmc.cumsum() 74 | #tmp_cmc = [x / (i + 1.) for i, x in enumerate(tmp_cmc)] 75 | y = np.arange(1, tmp_cmc.shape[0] + 1) * 1.0 76 | tmp_cmc = tmp_cmc / y 77 | tmp_cmc = np.asarray(tmp_cmc) * orig_cmc 78 | AP = tmp_cmc.sum() / num_rel 79 | all_AP.append(AP) 80 | 81 | assert num_valid_q > 0, "Error: all query identities do not appear in gallery" 82 | 83 | all_cmc = np.asarray(all_cmc).astype(np.float32) 84 | all_cmc = all_cmc.sum(0) / num_valid_q 85 | mAP = np.mean(all_AP) 86 | 87 | return all_cmc, mAP 88 | 89 | 90 | class R1_mAP_eval(): 91 | def __init__(self, num_query, max_rank=50, feat_norm=True, reranking=False): 92 | super(R1_mAP_eval, self).__init__() 93 | self.num_query = num_query 94 | self.max_rank = max_rank 95 | self.feat_norm = feat_norm 96 | self.reranking = reranking 97 | 98 | def reset(self): 99 | self.feats = [] 100 | self.pids = [] 101 | self.camids = [] 102 | 103 | def update(self, output): # called once for each batch 104 | feat, pid, camid = output 105 | self.feats.append(feat.cpu()) 106 | self.pids.extend(np.asarray(pid)) 107 | self.camids.extend(np.asarray(camid)) 108 | 109 | def compute(self): # called after each epoch 110 | feats = torch.cat(self.feats, dim=0) 111 | if self.feat_norm: 112 | print("The test feature is normalized") 113 | feats = torch.nn.functional.normalize(feats, dim=1, p=2) # along channel 114 | # query 115 | qf = feats[:self.num_query] 116 | q_pids = np.asarray(self.pids[:self.num_query]) 117 | q_camids = np.asarray(self.camids[:self.num_query]) 118 | # gallery 119 | gf = feats[self.num_query:] 120 | g_pids = np.asarray(self.pids[self.num_query:]) 121 | 122 | g_camids = np.asarray(self.camids[self.num_query:]) 123 | if self.reranking: 124 | print('=> Enter reranking') 125 | # distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3) 126 | distmat = re_ranking(qf, gf, k1=50, k2=15, lambda_value=0.3) 127 | 128 | else: 129 | print('=> Computing DistMat with euclidean_distance') 130 | distmat = euclidean_distance(qf, gf) 131 | cmc, mAP = eval_func(distmat, q_pids, g_pids, q_camids, g_camids) 132 | 133 | return cmc, mAP, distmat, self.pids, self.camids, qf, gf 134 | 135 | 136 | 137 | -------------------------------------------------------------------------------- /utils/reranking.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Fri, 25 May 2018 20:29:09 5 | 6 | 7 | """ 8 | 9 | """ 10 | CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017. 11 | url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf 12 | Matlab version: https://github.com/zhunzhong07/person-re-ranking 13 | """ 14 | 15 | """ 16 | API 17 | 18 | probFea: all feature vectors of the query set (torch tensor) 19 | probFea: all feature vectors of the gallery set (torch tensor) 20 | k1,k2,lambda: parameters, the original paper is (k1=20,k2=6,lambda=0.3) 21 | MemorySave: set to 'True' when using MemorySave mode 22 | Minibatch: avaliable when 'MemorySave' is 'True' 23 | """ 24 | 25 | import numpy as np 26 | import torch 27 | 28 | 29 | def re_ranking(probFea, galFea, k1, k2, lambda_value, local_distmat=None, only_local=False): 30 | # if feature vector is numpy, you should use 'torch.tensor' transform it to tensor 31 | query_num = probFea.size(0) 32 | all_num = query_num + galFea.size(0) 33 | if only_local: 34 | original_dist = local_distmat 35 | else: 36 | feat = torch.cat([probFea, galFea]) 37 | # print('using GPU to compute original distance') 38 | distmat = torch.pow(feat, 2).sum(dim=1, keepdim=True).expand(all_num, all_num) + \ 39 | torch.pow(feat, 2).sum(dim=1, keepdim=True).expand(all_num, all_num).t() 40 | distmat.addmm_(1, -2, feat, feat.t()) 41 | original_dist = distmat.cpu().numpy() 42 | del feat 43 | if not local_distmat is None: 44 | original_dist = original_dist + local_distmat 45 | gallery_num = original_dist.shape[0] 46 | original_dist = np.transpose(original_dist / np.max(original_dist, axis=0)) 47 | V = np.zeros_like(original_dist).astype(np.float16) 48 | initial_rank = np.argsort(original_dist).astype(np.int32) 49 | 50 | # print('starting re_ranking') 51 | for i in range(all_num): 52 | # k-reciprocal neighbors 53 | forward_k_neigh_index = initial_rank[i, :k1 + 1] 54 | backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1] 55 | fi = np.where(backward_k_neigh_index == i)[0] 56 | k_reciprocal_index = forward_k_neigh_index[fi] 57 | k_reciprocal_expansion_index = k_reciprocal_index 58 | for j in range(len(k_reciprocal_index)): 59 | candidate = k_reciprocal_index[j] 60 | candidate_forward_k_neigh_index = initial_rank[candidate, :int(np.around(k1 / 2)) + 1] 61 | candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index, 62 | :int(np.around(k1 / 2)) + 1] 63 | fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0] 64 | candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate] 65 | if len(np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)) > 2 / 3 * len( 66 | candidate_k_reciprocal_index): 67 | k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index, candidate_k_reciprocal_index) 68 | 69 | k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) 70 | weight = np.exp(-original_dist[i, k_reciprocal_expansion_index]) 71 | V[i, k_reciprocal_expansion_index] = weight / np.sum(weight) 72 | original_dist = original_dist[:query_num, ] 73 | if k2 != 1: 74 | V_qe = np.zeros_like(V, dtype=np.float16) 75 | for i in range(all_num): 76 | V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0) 77 | V = V_qe 78 | del V_qe 79 | del initial_rank 80 | invIndex = [] 81 | for i in range(gallery_num): 82 | invIndex.append(np.where(V[:, i] != 0)[0]) 83 | 84 | jaccard_dist = np.zeros_like(original_dist, dtype=np.float16) 85 | 86 | for i in range(query_num): 87 | temp_min = np.zeros(shape=[1, gallery_num], dtype=np.float16) 88 | indNonZero = np.where(V[i, :] != 0)[0] 89 | indImages = [invIndex[ind] for ind in indNonZero] 90 | for j in range(len(indNonZero)): 91 | temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum(V[i, indNonZero[j]], 92 | V[indImages[j], indNonZero[j]]) 93 | jaccard_dist[i] = 1 - temp_min / (2 - temp_min) 94 | 95 | final_dist = jaccard_dist * (1 - lambda_value) + original_dist * lambda_value 96 | del original_dist 97 | del V 98 | del jaccard_dist 99 | final_dist = final_dist[:query_num, query_num:] 100 | return final_dist 101 | 102 | --------------------------------------------------------------------------------