DOCSEARCHbyAlgolia
Submit Docs

TransReID

repository·main·Indexed 21 days ago

https://github.com/damo-cv/transreid

A Transformer-based framework for person and vehicle Object Re-Identification. It leverages transformer architectures and specialized modules such as Joint Part Matching (JPM) and Side Information Embedding (SIE) to achieve state-of-the-art performance. The repository includes support for datasets like Market-1501, MSMT17, DukeMTMC-reID, Occluded-Duke, VehicleID, and VeRi, with compatibility for ViT and DeiT pre-trained models.

Tokens
1.5K
Snippets
6
Records
7
Agent score
27%

What's inside TransReID

  1. Download pre-trained DeiT or ViT models

    main

    TransReID requires ImageNet pre-trained transformer models. Download one of the following:

    • ViT-Base: jx_vit_base_p16_224-80ecf9dd.pth
    • ViT-Small: vit_small_p16_224-15ec54c9.pth
    • DeiT-Small: deit_small_distilled_patch16_224-649709d9.pth
    • DeiT-Base: deit_base_distilled_patch16_224-df68dfff.pth

  2. Prepare datasets for TransReID

    main

    Create a data directory and organize your downloaded datasets (Market-1501, MSMT17, DukeMTMC-reID, Occluded-Duke, VehicleID, and VeRi) into the following structure:

    data
├── market1501
│   └── images ..
├── MSMT17
│   └── images ..
├── dukemtmcreid
│   └── images ..
├── Occluded_Duke
│   └── images ..
├── VehicleID_V1.0
│   └── images ..
└── VeRi
    └── images ..
    mkdir data

  3. Train TransReID using pre-defined configurations

    main

    TransReID provides several .yml configuration files for different datasets and model variants (Baseline, JPM, SIE, or full TransReID).

    Common Training Commands:

    DukeMTMC variants:

    • Baseline: configs/DukeMTMC/vit_base.yml
    • Baseline + JPM: configs/DukeMTMC/vit_jpm.yml
    • Baseline + SIE: configs/DukeMTMC/vit_sie.yml
    • TransReID (Baseline + SIE + JPM): configs/DukeMTMC/vit_transreid.yml
    • TransReID with stride [12, 12]: configs/DukeMTMC/vit_transreid_stride.yml

    Other Datasets:

    • MSMT17: configs/MSMT17/vit_transreid_stride.yml
    • OCC_Duke: configs/OCC_Duke/vit_transreid_stride.yml
    • Market: configs/Market/vit_transreid_stride.yml
    • VeRi: configs/VeRi/vit_transreid_stride.yml

    Distributed Training (VehicleID): Since VehicleID is large, use 4 V100 GPUs via torch.distributed.launch or the provided dist_train.sh script.

    # DukeMTMC transformer-based baseline
python train.py --config_file configs/DukeMTMC/vit_base.yml MODEL.DEVICE_ID "('0')"

# VehicleID (Distributed training with 4 GPUs)
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_transreid_stride.yml MODEL.DIST_TRAIN True

# Or use the shell script
Bash dist_train.sh

  4. Install TransReID and meet requirements

    main

    Install the necessary dependencies using pip.

    Recommended Environment:

    • Python >= 3.5
    • PyTorch >= 1.6 (required for torch.cuda.amp acceleration)
    • torchvision 0.7.0
    • timm 0.3.2
    • CUDA 10.1
    • Hardware: V100 (16GB or 32GB) is recommended for training and evaluation.

    Note on Memory: For person datasets with size 256x128, TransReID with stride occupies ~12GB GPU memory, while standard TransReID occupies ~7GB GPU memory.

    pip install -r requirements.txt

  5. Train TransReID with custom arguments

    main

    You can train the model using a single GPU by passing specific arguments to train.py.

    Arguments:

    • ${1}: Stride size for pure transformer (e.g., [16, 16], [14, 14], [12, 12]).
    • ${2}: True or False to use SIE with camera.
    • ${3}: True or False to use SIE with view.
    • ${4}: True or False to use JPM.
    • ${5}: Transformer type. Options: 'vit_base_patch16_224_TransReID', 'vit_small_patch16_224_TransReID', or 'deit_small_patch16_224_TransReID'.
    python train.py --config_file configs/transformer_base.yml MODEL.DEVICE_ID "('your device id')" MODEL.STRIDE_SIZE ${1} MODEL.SIE_CAMERA ${2} MODEL.SIE_VIEW ${3} MODEL.JPM ${4} MODEL.TRANSFORMER_TYPE ${5} OUTPUT_DIR ${OUTPUT_DIR} DATASETS.NAMES "('your dataset name')"

  6. Evaluate TransReID models

    main

    To evaluate a trained model, use test.py and provide the corresponding configuration file and the path to your trained checkpoint (TEST.WEIGHT).

    Example for DukeMTMC:

    python test.py --config_file configs/DukeMTMC/vit_transreid_stride.yml MODEL.DEVICE_ID "('0')" TEST.WEIGHT '../logs/duke_vit_transreid_stride/transformer_120.pth'
    python test.py --config_file 'choose which config to test' MODEL.DEVICE_ID "('your device id')" TEST.WEIGHT "('your path of trained checkpoints')"

  7. Download trained TransReID models and logs

    main

    Pre-trained models and training logs for various datasets (MSMT17, Market, Duke, OCC_Duke, VeRi, and VehicleID) are available via Google Drive. The models include different architectures:

    • Baseline (ViT)
    • TransReID (ViT)*
    • TransReID (DeiT)*

    Note that performance metrics (mAP | R1) in this repository may differ slightly from the original paper due to code reorganization.

    | Model | MSMT17 (mAP|R1) | Market (mAP|R1) | Duke (mAP|R1) | OCC_Duke (mAP|R1) | VeRi (mAP|R1) | VehicleID (R1|R5) |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| Baseline(ViT) | 61.8 | 81.8 | 87.1 | 94.6 | 79.6 | 89.0 | 53.8 | 61.1 | 79.0 | 96.6 | 83.5 | 96.7 |
| TransReID*(ViT) | 67.8 | 85.3 | 89.0 | 95.1 | 82.2 | 90.7 | 59.5 | 67.4 | 82.1 | 97.4 | 85.2 | 97.4 |
| TransReID*(DeiT) | 66.3 | 84.0 | 88.5 | 95.1 | 81.9 | 90.7 | 57.7 | 65.2 | 82.4 | 97.1 | 86.0 | 97.6 |