DOCSEARCHbyAlgolia
Submit Docs

StreamPETR Documentation

repository·main·Indexed 21 days ago

https://github.com/exiawsh/streampetr

Official implementation of an ICCV 2023 paper on object-centric temporal modeling for efficient multi-view 3D object detection and tracking. Features include support for StreamPETR, PETR, and Focal-PETR architectures, Flash and Deformable attention, and TensorRT inference. Includes guides for environment setup using Conda and PyTorch 1.9.0, nuScenes dataset preparation, and training modes such as Sliding Window and Streaming Video.

Tokens
4.2K
Snippets
11
Records
19
Agent score
74%

What's inside StreamPETR

  1. StreamPETR Supported Features

    main

    StreamPETR supports the following core functionalities:

    • Core Architectures: StreamPETR (including PETR and Focal-PETR).
    • Attention Mechanisms: Flash attention and Deformable attention (via RepDETR3D).
    • Training Techniques: Sliding window training and efficient training for streaming video.
    • Inference & Tasks: TensorRT inference and 3D object tracking.

  2. Getting Started with StreamPETR

    main

    To use StreamPETR, follow these three primary steps in order:

    1. Environment Setup: Configure your development environment.
    2. Data Preparation: Prepare the necessary datasets (e.g., NuScenes).
    3. Training and Inference: Execute model training or run inference on data.

    Detailed instructions for each step are located in the project's docs/ directory.

  3. Transform EVA02 pretrained weights for ViT-Large

    main

    To use EVA02 pretrained weights (such as Object365 or ImageNet weights) with the StreamPETR ViT-Large backbone, you must remap the keys in the state dictionary. Specifically, backbone.net keys must be remapped to img_backbone. and backbone.simfp keys must be remapped to img_backbone.adapter..

    import torch

pretrain_dict = torch.load('ckpts/eva02_L_coco_det_sys_o365.pth', map_location=torch.device('cpu'))
pretrain_dict = pretrain_dict["model"]
print(pretrain_dict.keys())
remapped_dict = {}
for k,v in pretrain_dict.items():
    if "backbone.net" in k:
        remapped_dict[k.replace("backbone.net.", "img_backbone.")] = v
    if "backbone.simfp" in k:
        remapped_dict[k.replace("backbone.", "img_backbone.adapter.")] = v
torch.save(remapped_dict,'ckpts/eva02_L_coco_det_sys_o365_remapped.pth')

  4. Evaluate StreamPETR detection and tracking models

    main

    Detection Evaluation

    Evaluate a detection model using tools/dist_test.sh. You must provide the config file, the path to the checkpoint (.pth), the number of GPUs, and the evaluation mode (e.g., --eval bbox).

    tools/dist_test.sh projects/configs/StreamPETR/stream_petr_vov_flash_800_bs2_seq_24e.py work_dirs/stream_petr_vov_flash_800_bs2_seq_24e/latest.pth 8 --eval bbox

    Tracking Evaluation

    Evaluate the tracking model using the nusc_tracking/pub_test script. Requires the NuScenes version, the path to the results JSON checkpoint, and the NuScenes data root.

    python nusc_tracking/pub_test --version v1.0-trainval --checkpoint {PATH_RESULTS.JSON} --data_root {PATH_NUSCENES}

  5. Install StreamPETR via Conda and Pip

    main

    Follow these steps to set up a clean environment for StreamPETR. This process involves creating a Conda environment, installing PyTorch with CUDA support, optionally installing flash-attn for performance, and installing the required OpenMMLab dependencies (mmdet3d, mmdet, and mmsegmentation).

    # 1. Create and activate conda environment
conda create -n streampetr python=3.8 -y
conda activate streampetr

# 2. Install PyTorch and torchvision
pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

# 3. Install flash-attn (optional)
pip install flash-attn==0.2.2

# 4. Clone StreamPETR
git clone https://github.com/exiawsh/StreamPETR

# 5. Install mmdet3d dependencies
pip install mmcv-full==1.6.0 -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html
pip install mmdet==2.28.2
pip install mmsegmentation==0.30.0

# 6. Install mmdet3d from source
cd ./StreamPETR
git clone https://github.com/open-mmlab/mmdetection3d.git
cd mmdetection3d
git checkout v1.0.0rc6 
pip install -e .

  6. Training recipes and performance optimization tricks

    main

    The following configuration adjustments can be used to boost StreamPETR performance:

    • Loss Weights: For SOTA results on sparse query designs, change the x,y weight from 1.0 to 2.0 in the bounding box regression loss/Hungarian matching.
    • Backbone Learning Rate:
      • For large backbones (e.g., VIT-Base) or 2D pretrained (e.g., R50-Nuimage), use 0.1.
      • For small IN1k-pretrained models (e.g., R50-IN1k), use 0.25 or 0.5.
    • Stability for Small Models: For small IN1k models, enable Sync-BN by setting SyncBN=True and updating the norm config:
      norm_cfg=dict(type='BN2d', requires_grad=True),
norm_eval=False
    • Inference Speed: Setting feedforward_channels for the Transformer to a smaller value (e.g., 512) can improve inference speed with minimal accuracy impact.
    • Flash Attention Compatibility: If your device does not support Flash attention, use the following attention config:
      dict(type='PETRMultiheadAttention', embed_dims=256, num_heads=8, dropout=0.1, fp16=True)
    • Learning Rate Scaling: Adjust the learning rate based on the total batch size (Num_gpus * Batch_size):
      • 8: 2e-4
      • 16: 4e-4
      • 32: 6e-4

  7. Estimate StreamPETR inference speed

    main

    To benchmark latency (including data-processing, network FP32, and post-processing), use tools/benchmark.py. Note that the workers_per_gpu setting in your config will affect the measured speed because data processing time is included in the benchmark.

    python tools/benchmark.py projects/configs/test_speed/stream_petr_r50_704_bs2_seq_428q_nui_speed_test.py

  8. Visualize 3D object detection results

    main

    Generate results JSON

    To generate only the results JSON file without full evaluation, use the --format-only flag with dist_test.sh:

    ./tools/dist_test.sh projects/configs/StreamPETR/stream_petr_vov_flash_800_bs2_seq_24e.py work_dirs/stream_petr_vov_flash_800_bs2_seq_24e/latest.pth 8 --format-only

    Visualize 3D detections

    Use tools/visualize.py to visualize detections. Note: You must manually edit the results_nusc.json path inside the tools/visualize.py file before running.

    python3 tools/visualize.py

  10. Set up pretrained weights directory

    main

    Create a ckpts directory in the StreamPETR root folder to store pretrained weights.

    Available weights include:

    cd /path/to/StreamPETR
mkdir ckpts

  11. Train StreamPETR models

    main

    You can train StreamPETR using the tools/dist_train.sh script. The project provides two training modes via configuration files:

    1. Sliding Window: Higher accuracy (as reported in the paper) but consumes significantly more GPU memory and training time.
    2. Streaming Video: Follows the SOLOFusion approach. It is much faster (saves approximately 4x training hours) but converges more slowly (e.g., 90 epochs of streaming video is roughly equivalent to 60 epochs of sliding window).

    To start training with 8 GPUs using a streaming video config:

    tools/dist_train.sh projects/configs/StreamPETR/stream_petr_r50_flash_704_bs2_seq_24e.py 8 --work-dir work_dirs/stream_petr_r50_flash_704_bs2_seq_24e/

  12. Create nuScenes 2D temporal info files

    main

    StreamPETR requires modified data preparation that includes 2D annotations and temporal information. You can generate these using the tools/create_data_nusc.py script. This process will produce nuscenes2d_temporal_infos_{train,val}.pkl files.

    Alternatively, you can download the pre-processed .pkl files directly:

    python tools/create_data_nusc.py --root-path ./data/nuscenes --out-dir ./data/nuscenes --extra-tag nuscenes2d --version v1.0