daniel-code/TubeViT

An unofficial implementation of TubeViT in "Rethinking Video ViTs: Sparse Video Tubes for Joint Image and Video Learning"

TubeViT

An unofficial implementation of TubeViT in "Rethinking Video ViTs: Sparse Video Tubes for Joint Image and Video Learning"

Spec.

• Fixed Positional embedding
• Sparse Tube Construction
  • Multi-Tube
  • Interpolated Kernels
  • Space To Depth
  • config of tubes
• pipeline
  • training
  • evaluating
  • inference

Usage

This project is based on torch==1.13.1 and pytorch-lightning

Setup

1. Install requirements

   pip install -r requirements.txt
   

2. Download UFC101 dataset

Convert ViT pre-trained weight

Use convert_vit_weight.py to convert torch ViT pre-trained weight to TubeVit.

python scripts/convert_vit_weight.py --help                                                                              ✔ 
Usage: convert_vit_weight.py [OPTIONS]

Options:
  -nc, --num-classes INTEGER      num of classes of dataset.
  -f, --frames-per-clip INTEGER   frame per clip.
  -v, --video-size <INTEGER INTEGER>...
                                  frame per clip.
  -o, --output-path PATH          output model weight name.
  --help                          Show this message and exit.

Example

Convert ImageNet pre-trained weight to UCF101. --num-classes is 101 by default.

python scripts/convert_vit_weight.py

Train

Current train.py only train on pytorch UCF101 dataset. Change the dataset if needed.

--dataset-root and --annotation-path is based on torchvision.datasets.UCF101

python scripts/train.py --help

Usage: train.py [OPTIONS]

Options:
  -r, --dataset-root PATH         path to dataset.  [required]
  -a, --annotation-path PATH      path to dataset.  [required]
  -nc, --num-classes INTEGER      num of classes of dataset.
  -b, --batch-size INTEGER        batch size.
  -f, --frames-per-clip INTEGER   frame per clip.
  -v, --video-size <INTEGER INTEGER>...
                                  frame per clip.
  --max-epochs INTEGER            max epochs.
  --num-workers INTEGER
  --fast-dev-run
  --seed INTEGER                  random seed.
  --preview-video                 Show input video
  --help                          Show this message and exit.

Examples

python scripts/train.py -r path/to/dataset -a path/to/annotation

Evaluation

python scripts/evaluate.py --help

Usage: evaluate.py [OPTIONS]

Options:
  -r, --dataset-root PATH         path to dataset.  [required]
  -m, --model-path PATH           path to model weight.  [required]
  -a, --annotation-path PATH      path to dataset.  [required]
  --label-path PATH               path to classInd.txt.  [required]
  -nc, --num-classes INTEGER      num of classes of dataset.
  -b, --batch-size INTEGER        batch size.
  -f, --frames-per-clip INTEGER   frame per clip.
  -v, --video-size <INTEGER INTEGER>...
                                  frame per clip.
  --num-workers INTEGER
  --seed INTEGER                  random seed.
  --verbose                       Show input video
  --help                          Show this message and exit.

Examples

python scripts/evaluate.py -r path/to/dataset -a path/to/annotation

Model Architecture

Positional embedding