TimeSformer

repository·main·Indexed 23 days ago

https://github.com/facebookresearch/timesformer

An official PyTorch implementation of a space-time attention mechanism for video understanding and classification. The framework supports various self-attention schemes, including divided space-time, space-only, and joint space-time attention, with pretrained models available for K400, K600, SSv2, and HowTo100M. It includes support for high-resolution (TimeSformer-HR) and long-clip (TimeSformer-L) variants, as well as distributed training via Submitit.

Tokens
2.4K
Snippets
8
Records
13
Agent score
74%

What's inside TimeSformer

  1. Train TimeSformer-HR and TimeSformer-L variants

    main

    For more powerful variants, use the following configurations. Note that these require GPUs with approximately 32GB of memory:

    • TimeSformer-HR (16-frame clips, 448x448 resolution): Use configs/Kinetics/TimeSformer_divST_16x16_448.yaml.
    • TimeSformer-L (96-frame clips, 224x224 resolution): Use configs/Kinetics/TimeSformer_divST_96x4_224.yaml.
  2. Perform inference (testing)

    main

    To run inference/testing, use tools/run_net.py with a testing configuration file. You must set TRAIN.ENABLE to False and provide the path to your checkpoint via TEST.CHECKPOINT_FILE_PATH.

    python tools/run_net.py \
      --cfg configs/Kinetics/TimeSformer_divST_8x32_224_TEST.yaml \
      DATA.PATH_TO_DATA_DIR path_to_your_dataset \
      TEST.CHECKPOINT_FILE_PATH path_to_your_checkpoint \
      TRAIN.ENABLE False
  3. Prepare the Kinetics dataset

    main

    To use the Kinetics dataset with TimeSformer, follow these steps:

    1. Download the videos from the Kinetics dataset repository.
    2. Resize all videos so that the short edge is 256 pixels.
    3. Create CSV files for the training, validation, and testing sets named train.csv, val.csv, and test.csv respectively.

    The CSV files must follow this format:

    path_to_video_1 label_1
    path_to_video_2 label_2
    path_to_video_3 label_3
    ...
    path_to_video_N label_N
  4. Train the default TimeSformer

    main

    To train the default TimeSformer (divided space-time attention, 8-frame clips, 224x224 resolution), use tools/run_net.py with a configuration file.

    You can specify the dataset path via the command line using DATA.PATH_TO_DATA_DIR or by adding it directly to the YAML config file.

    Example Command:

    python tools/run_net.py \
      --cfg configs/Kinetics/TimeSformer_divST_8x32_224.yaml \
      DATA.PATH_TO_DATA_DIR path_to_your_dataset \
      NUM_GPUS 8 \
      TRAIN.BATCH_SIZE 8
  5. Distributed training via Submitit

    main

    For multi-node distributed training, install submitit and use tools/submit.py.

    Example: Training on Kinetics using 4 nodes (8 GPUs each):

    python tools/submit.py --cfg configs/Kinetics/TimeSformer_divST_8x32_224.yaml --job_dir /your/job/dir/${JOB_NAME}/ --num_shards 4 --name ${JOB_NAME} --use_volta32
    pip install submitit
    
    python tools/submit.py --cfg configs/Kinetics/TimeSformer_divST_8x32_224.yaml --job_dir /your/job/dir/${JOB_NAME}/ --num_shards 4 --name ${JOB_NAME} --use_volta32
  6. Experiment with different self-attention schemes

    main

    You can switch between different attention mechanisms by selecting the appropriate configuration file in tools/run_net.py:

    • Divided Space-Time Attention: Use configs/Kinetics/TimeSformer_divST_...yaml (Default).
    • Space-Only Attention: Use configs/Kinetics/TimeSformer_spaceOnly_...yaml.
    • Joint Space-Time Attention: Use configs/Kinetics/TimeSformer_jointST_...yaml.
  7. Prepare the Something-Something V2 dataset

    main

    To use the Something-Something V2 dataset, follow these steps:

    1. Download the dataset and annotations from the official provider.
    2. Download the official frame lists for training and validation:
    3. Extract frames from the videos at 30 FPS using ffmpeg. Ensure the resulting frame structure matches the frame lists.
    4. Organize your files:
      • Place all annotation JSON files and the frame lists in a single folder.
      • Set the configuration key DATA.PATH_TO_DATA_DIR to the path of that folder.
      • Set the configuration key DATA.PATH_PREFIX to the path of the folder containing the extracted frames.
    ffmpeg -i "${video}" -r 30 -q:v 1 "${out_name}"
  8. Install TimeSformer

    main

    To install TimeSformer, first create and activate a conda environment with Python 3.7. Then, install the required dependencies via pip or conda, and finally build the codebase using setup.py.

    Dependencies:

    • torchvision
    • fvcore (from git+https://github.com/facebookresearch/fvcore)
    • simplejson
    • einops
    • timm
    • av (PyAV)
    • psutil
    • scikit-learn
    • opencv-python
    • tensorboard

    Note: The environment was developed on Ubuntu 20.04.

    conda create -n timesformer python=3.7 -y
    source activate timesformer
    
    # Install dependencies (example list)
    pip install torchvision
    pip install 'git+https://github.com/facebookresearch/fvcore'
    pip install simplejson einops timm psutil scikit-learn opencv-python tensorboard
    conda install av -c conda-forge
    
    # Build the codebase
    git clone https://github.com/facebookresearch/TimeSformer
    cd TimeSformer
    python setup.py build develop
  9. Finetune TimeSformer from a checkpoint

    main

    To finetune a model from an existing PyTorch checkpoint, add the following flags to your training command or include them in your YAML config:

    • TRAIN.CHECKPOINT_FILE_PATH: path to your .pyth checkpoint.
    • TRAIN.FINETUNE: set to True.
  10. Configure training for different GPU counts

    main

    To use a different number of GPUs, you must modify the .yaml configuration files in the configs/ directory.

    Update the following keys:

    • NUM_GPUS
    • TRAIN.BATCH_SIZE (must be $\ge$ NUM_GPUS)
    • TEST.BATCH_SIZE
    • DATA_LOADER.NUM_WORKERS

    A sample configuration for a 4 GPU setup is provided at configs/Kinetics/TimeSformer_divST_8x32_224_4gpus.yaml.

  11. Load pretrained TimeSformer models

    main

    You can load pretrained TimeSformer models (available for K400, K600, SSv2, and HowTo100M) using the TimeSformer class from timesformer.models.vit.

    When initializing, specify the img_size, num_classes, num_frames, and attention_type. Use the pretrained_model argument to provide the path to the downloaded .pyth file.

    Input Shape: The model expects a video tensor of shape (batch x channels x frames x height x width).

    import torch
    from timesformer.models.vit import TimeSformer
    
    model = TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type='divided_space_time',  pretrained_model='/path/to/pretrained/model.pyth')
    
    dummy_video = torch.randn(2, 3, 8, 224, 224) # (batch x channels x frames x height x width)
    
    pred = model(dummy_video,)
  12. Initialize and run inference with TimeSformer

    main

    To use a pre-trained TimeSformer model for inference, import the TimeSformer class from timesformer.models.vit. You must specify the img_size, num_classes, num_frames, and attention_type. To load weights, provide the path to a .pyth model file via the pretrained_model argument.

    Input tensors should be in the shape (batch, channels, frames, height, width).

    from pathlib import Path
    import torch
    from timesformer.models.vit import TimeSformer
    
    # Define path to pre-trained weights
    model_file = Path.home()/'TimeSformer/models/TimeSformer_divST_8x32_224_K600.pyth'
    
    # Initialize model
    model = TimeSformer(
        img_size=224, 
        num_classes=600, 
        num_frames=8, 
        attention_type='divided_space_time',  
        pretrained_model=str(model_file)
    )
    
    # Prepare dummy video: (batch x channels x frames x height x width)
    dummy_video = torch.randn(2, 3, 8, 224, 224)
    
    # Perform inference
    pred = model(dummy_video)
    
    # Output shape will be (batch, num_classes)
    assert pred.shape == (2, 600)