uzh-rpg/ssms_event_cameras

Official PyTorch implementation of the CVPR 2024 paper: State Space Models for Event Cameras.

State Space Models for Event Cameras

This is the official PyTorch implementation of the CVPR 2024 paper State Space Models for Event Cameras.

✅ Updates

• April. 19th, 2024: The code along with the best checkpoints is released! The poster and video will be released shortly before CVPR 2024.

Citation

If you find this work and/or code useful, please cite our paper:

@InProceedings{Zubic_2024_CVPR,
  author  = {Zubi\'c, Nikola and Gehrig, Mathias and Scaramuzza, Davide},
  title   = {State Space Models for Event Cameras},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year    = {2024},
}

SSM-ViT

• S5 model used in our SSM-ViT pipeline can be seen here.
• In particular, S5 is used instead of RNN in a 4-stage hierarchical ViT backbone, and its forward function is exposed here. What is nice about this approach is that we do not need a 'for' loop over sequence dimension, but instead we employ a parallel scanning algorithm. This model assumes that a hidden state is being carried over.
• For a model that is standalone, and can be used for any sequence modeling problem, one does not use by default this formulation where we carry on the hidden state. The implementation is the same as the original JAX implementation and can be downloaded in zip format from ssms_event_cameras/RVT/models/s5.zip.

Installation

Conda

We highly recommend using Mambaforge to reduce the installation time.

conda create -y -n events_signals python=3.11
conda activate events_signals
conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install lightning wandb pandas plotly opencv-python tabulate pycocotools bbox-visualizer StrEnum hydra-core einops torchdata tqdm numba h5py hdf5plugin lovely-tensors tensorboardX pykeops scikit-learn          

Required Data

To evaluate or train the S5-ViT model, you will need to download the required preprocessed datasets:

	1 Mpx	Gen1
pre-processed dataset	download	download
crc32	c5ec7c38	5acab6f3

You may also pre-process the dataset yourself by following the instructions.

Pre-trained Checkpoints

1 Mpx

	S5-ViT-Base	S5-ViT-Small
pre-trained checkpoint	download	download

Gen1

	S5-ViT-Base	S5-ViT-Small
pre-trained checkpoint	download	download

Evaluation

• Evaluation scripts with concrete parameters that we trained our models can be seen here.

• Set DATA_DIR as the path to either the 1 Mpx or Gen1 dataset directory

• Set CKPT_PATH to the path of the correct checkpoint matching the choice of the model and dataset

• Set

  • MDL_CFG=base or
  • MDL_CFG=small

  to load either the base or small model configuration.

• Set GPU_ID to the PCI BUS ID of the GPU that you want to use. e.g. GPU_ID=0. Only a single GPU is supported for evaluation

1 Mpx

python RVT/validation.py dataset=gen4 dataset.path=${DATA_DIR} checkpoint=${CKPT_PATH} \
use_test_set=1 hardware.gpus=${GPU_ID} +experiment/gen4="${MDL_CFG}.yaml" \
batch_size.eval=12 model.postprocess.confidence_threshold=0.001

Gen1

python RVT/validation.py dataset=gen1 dataset.path=${DATA_DIR} checkpoint=${CKPT_PATH} \
use_test_set=1 hardware.gpus=${GPU_ID} +experiment/gen1="${MDL_CFG}.yaml" \
batch_size.eval=8 model.postprocess.confidence_threshold=0.001

We set the same batch size for the evaluation and training: 12 for the 1 Mpx dataset, and 8 for the Gen1 dataset.

Evaluation results

Evaluation should give the same results as shown below:

• 47.7 and 47.8 mAP on Gen1 and 1 Mpx datasets for the base model, and
• 46.6 and 46.5 mAP on Gen1 and 1 Mpx datasets for the small model.

Training

• Set DATA_DIR as the path to either the 1 Mpx or Gen1 dataset directory

• Set

  • MDL_CFG=base or
  • MDL_CFG=small

  to load either the base or the small configuration.

• Set GPU_IDS to the PCI BUS IDs of the GPUs that you want to use. e.g. GPU_IDS=[0,1] for using GPU 0 and 1. Using a list of IDS will enable single-node multi-GPU training. Pay attention to the batch size which is defined per GPU.

• Set BATCH_SIZE_PER_GPU such that the effective batch size is matching the parameters below. The effective batch size is (batch size per GPU)*(number of GPUs).

• If you would like to change the effective batch size, we found the following learning rate scaling to work well for all models on both datasets:

  lr = 2e-4 * sqrt(effective_batch_size/8).

• The training code uses W&B for logging during the training. Hence, we assume that you have a W&B account.

  • The training script below will create a new project called ssms_event_cameras. Adapt the project name and group name if necessary.

1 Mpx

• The effective batch size for the 1 Mpx training is 12.
• For training the model on 1 Mpx dataset, we need 2x A100 80 GB GPUs and we use 12 workers per GPU for training and 4 workers per GPU for evaluation:

GPU_IDS=[0,1]
BATCH_SIZE_PER_GPU=6
TRAIN_WORKERS_PER_GPU=12
EVAL_WORKERS_PER_GPU=4
python RVT/train.py model=rnndet dataset=gen4 dataset.path=${DATA_DIR} wandb.project_name=ssms_event_cameras \
wandb.group_name=1mpx +experiment/gen4="${MDL_CFG}.yaml" hardware.gpus=${GPU_IDS} \
batch_size.train=${BATCH_SIZE_PER_GPU} batch_size.eval=${BATCH_SIZE_PER_GPU} \
hardware.num_workers.train=${TRAIN_WORKERS_PER_GPU} hardware.num_workers.eval=${EVAL_WORKERS_PER_GPU}

If you for example want to execute the training on 4 GPUs simply adapt GPU_IDS and BATCH_SIZE_PER_GPU accordingly:

GPU_IDS=[0,1,2,3]
BATCH_SIZE_PER_GPU=3

Gen1

• The effective batch size for the Gen1 training is 8.
• For training the model on the Gen1 dataset, we need 1x A100 80 GPU using 24 workers for training and 8 workers for evaluation:

GPU_IDS=0
BATCH_SIZE_PER_GPU=8
TRAIN_WORKERS_PER_GPU=24
EVAL_WORKERS_PER_GPU=8
python RVT/train.py model=rnndet dataset=gen1 dataset.path=${DATA_DIR} wandb.project_name=ssms_event_cameras \
wandb.group_name=gen1 +experiment/gen1="${MDL_CFG}.yaml" hardware.gpus=${GPU_IDS} \
batch_size.train=${BATCH_SIZE_PER_GPU} batch_size.eval=${BATCH_SIZE_PER_GPU} \
hardware.num_workers.train=${TRAIN_WORKERS_PER_GPU} hardware.num_workers.eval=${EVAL_WORKERS_PER_GPU}

Code Acknowledgments

This project has used code from the following projects:

• RVT - Recurrent Vision Transformers for Object Detection with Event Cameras in PyTorch
• S4 - Structured State Spaces for Sequence Modeling, in particular S4 and S4D models in PyTorch
• S5 - Simplified State Space Layers for Sequence Modeling in JAX

ssms_event_cameras's Issues

Looking forward to code release

Thanks for sharing this work! Looking forward to the code release~

evaluation time question

Hello! I have been interested in RVT for a long time . I think you have done a great job! I have a question and I hope you can help me answer it. After I reproduce RVT, I have the following reasoning time:
Loading and preparing results DONE (t=0.41s) creating index Index created!
Running per image evaluation Evaluate annotation type * bbox * DONE (t=7.23s)
Accumulating evaluation results DONE (t=2.31s)
May I ask how to infer the time velocity in ms, just like 7.81 ms in your paper on the gen1 dataset.
I am using a 3090TIGPU with a batch_size of 1, so what should be the speed in ms?

Testing in a higher frequency

Hi @NikolaZubic

Thanks for your nice work and opening source!
I have a question when testing in a higher frequency, which parameters we should change?

I have found two places about the 'step_scale':

class S5SSM(torch.nn.Module):
    def __init__(
        self,
        lambdaInit: torch.Tensor,
        V: torch.Tensor,
        Vinv: torch.Tensor,
        h: int,
        p: int,
        dt_min: float,
        dt_max: float,
        liquid: bool = False,
        factor_rank: Optional[int] = None,
        discretization: Literal["zoh", "bilinear"] = "bilinear",
        bcInit: Initialization = "factorized",
        degree: int = 1,
        bidir: bool = False,
        step_scale: float = 1.0,
        bandlimit: Optional[float] = None,
    ):

and

    def forward(self, signal, prev_state, step_scale: float | torch.Tensor = 1.0): 

in S5SSM moudle.

If I want two testing in frequecy 200Hz (training in frequency 20), how should I adjust the parameter?
Should I change the above step_scale from 1.0 to 0.1 both?