botaoye / OSTrack

[ECCV 2022] Joint Feature Learning and Relation Modeling for Tracking: A One-Stream Framework
MIT License
393 stars 50 forks source link
tracking transformer

OSTrack

The official implementation for the ECCV 2022 paper Joint Feature Learning and Relation Modeling for Tracking: A One-Stream Framework.

[Models][Raw Results][Training logs]

Framework

News

[Dec. 12, 2022]

[Oct. 28, 2022]

Highlights

:star2: New One-stream Tracking Framework

OSTrack is a simple, neat, high-performance one-stream tracking framework for joint feature learning and relational modeling based on self-attention operators. Without any additional temporal information, OSTrack achieves SOTA performance on multiple benchmarks. OSTrack can serve as a strong baseline for further research.

Tracker GOT-10K (AO) LaSOT (AUC) TrackingNet (AUC) UAV123(AUC)
OSTrack-384 73.7 71.1 83.9 70.7
OSTrack-256 71.0 69.1 83.1 68.3

:star2: Fast Training

OSTrack-256 can be trained in ~24 hours with 4*V100 (16GB of memory per GPU), which is much faster than recent SOTA transformer-based trackers. The fast training speed comes from:

  1. While previous Siamese-style trackers required separate feeding of the template and search region into the backbone at each iteration of training, OSTrack directly combines the template and search region. The tight and highly parallelized structure results in improved training and inference speed.

  2. The proposed early candidate elimination (ECE) module significantly reduces memory and time consumption.

  3. Pretrained Transformer weights enable faster convergence.

:star2: Good performance-speed trade-off

speed_vs_performance

Install the environment

Option1: Use the Anaconda (CUDA 10.2)

conda create -n ostrack python=3.8
conda activate ostrack
bash install.sh

Option2: Use the Anaconda (CUDA 11.3)

conda env create -f ostrack_cuda113_env.yaml

Option3: Use the docker file

We provide the full docker file here.

Set project paths

Run the following command to set paths for this project

python tracking/create_default_local_file.py --workspace_dir . --data_dir ./data --save_dir ./output

After running this command, you can also modify paths by editing these two files

lib/train/admin/local.py  # paths about training
lib/test/evaluation/local.py  # paths about testing

Data Preparation

Put the tracking datasets in ./data. It should look like this:

   ${PROJECT_ROOT}
    -- data
        -- lasot
            |-- airplane
            |-- basketball
            |-- bear
            ...
        -- got10k
            |-- test
            |-- train
            |-- val
        -- coco
            |-- annotations
            |-- images
        -- trackingnet
            |-- TRAIN_0
            |-- TRAIN_1
            ...
            |-- TRAIN_11
            |-- TEST

Training

Download pre-trained MAE ViT-Base weights and put it under $PROJECT_ROOT$/pretrained_models (different pretrained models can also be used, see MAE for more details).

python tracking/train.py --script ostrack --config vitb_256_mae_ce_32x4_ep300 --save_dir ./output --mode multiple --nproc_per_node 4 --use_wandb 1

Replace --config with the desired model config under experiments/ostrack. We use wandb to record detailed training logs, in case you don't want to use wandb, set --use_wandb 0.

Evaluation

Download the model weights from Google Drive

Put the downloaded weights on $PROJECT_ROOT$/output/checkpoints/train/ostrack

Change the corresponding values of lib/test/evaluation/local.py to the actual benchmark saving paths

Some testing examples:

Visualization or Debug

Visdom is used for visualization.

  1. Alive visdom in the server by running visdom:

  2. Simply set --debug 1 during inference for visualization, e.g.:

    python tracking/test.py ostrack vitb_384_mae_ce_32x4_ep300 --dataset vot22 --threads 1 --num_gpus 1 --debug 1
  3. Open http://localhost:8097 in your browser (remember to change the IP address and port according to the actual situation).

  4. Then you can visualize the candidate elimination process.

ECE_vis

Test FLOPs, and Speed

Note: The speeds reported in our paper were tested on a single RTX2080Ti GPU.

# Profiling vitb_256_mae_ce_32x4_ep300
python tracking/profile_model.py --script ostrack --config vitb_256_mae_ce_32x4_ep300
# Profiling vitb_384_mae_ce_32x4_ep300
python tracking/profile_model.py --script ostrack --config vitb_384_mae_ce_32x4_ep300

Acknowledgments

Citation

If our work is useful for your research, please consider citing:

@inproceedings{ye2022ostrack,
  title={Joint Feature Learning and Relation Modeling for Tracking: A One-Stream Framework},
  author={Ye, Botao and Chang, Hong and Ma, Bingpeng and Shan, Shiguang and Chen, Xilin},
  booktitle={ECCV},
  year={2022}
}