liming-ai / tsm

Reproduce TSM for action recognition

Introduction

This is the code for TSM: Temporal Shift Module for Efficient Video Understanding on CVPR 2019 and Temporal Segment Networks: Towards Good Practices for Deep Action Recognition on ECCV 2016.

We only support for UCF101 rawframe dataset, but you can modify the dataset.py to add more dataset.

This repository is designed to help beginners better understand video classification, about how to write the latest models and understand the data pipelines.

So we added as much detailed code comments as possible, the basic code is from official TSM implementation, we only rewrite the code for brevity and easier understanding.

Prerequisites

You should install all the required dependencies:

• PyTorch 1.0 or higher
• TensorboardX
• tqdm
• scikit-learn

Data Preparation

Here we only write code for rawframes of UCF101, please see mmaction2 to prepare these frames.

The folder structure should be like this:

TSM
├── backbones
├── dataset.py
├── logger.py
├── opts.py
├── README.md
├── test.py
├── train.py
├── transforms.py
├── tsn.py
├── utils.py
├── data
│   ├── ucf101
│   │   ├── ucf101_{train,val}_split_{1,2,3}_rawframes.txt
│   │   ├── ucf101_{train,val}_split_{1,2,3}_videos.txt
│   │   ├── annotations
│   │   ├── rawframes
│   │   │   ├── ApplyEyeMakeup
│   │   │   │   ├── v_ApplyEyeMakeup_g01_c01
│   │   │   │   │   ├── img_00001.jpg
│   │   │   │   │   ├── img_00002.jpg
│   │   │   │   │   ├── ...
│   │   │   │   │   ├── flow_x_00001.jpg
│   │   │   │   │   ├── flow_x_00002.jpg
│   │   │   │   │   ├── ...
│   │   │   │   │   ├── flow_y_00001.jpg
│   │   │   │   │   ├── flow_y_00002.jpg
│   │   │   ├── ...
│   │   │   ├── YoYo
│   │   │   │   ├── v_YoYo_g01_c01
│   │   │   │   ├── ...
│   │   │   │   ├── v_YoYo_g25_c05

Use

Before use the code, you should create a folder to store all the training logs and checkpoints:

mkdir output_dir

You can find all the needed parameters in ots.py, set the parameters you want and run the code for training directly:

python train.py

or run the code for testing:

python test.py

You can set these parameters in terminal too:

python train.py --backbone "resnet50_tsm" --batch_size 32 --epochs 25

We also support pretrained models from TSM pretrained models, just add code in terminal like this:

python train.py --load_from_github True --state_dict_path path_to_pretrain_model

Please read the code in opts.py for more details.

Analysis

You can find the training logs in output_dir, and use tensorboard to analyze the results by running:

tensorboard --logdir output_dir