78 lines
3.2 KiB
Markdown
78 lines
3.2 KiB
Markdown
# Temporal Shift Module
|
|
|
|
TSM-R50 from "TSM: Temporal Shift Module for Efficient Video Understanding" <https://arxiv.org/abs/1811.08383>
|
|
|
|
TSM is a widely used Action Recognition model. This TensorRT implementation is tested with TensorRT 5.1 and TensorRT 7.2.
|
|
|
|
For the PyTorch implementation, you can refer to [open-mmlab/mmaction2](https://github.com/open-mmlab/mmaction2) or [mit-han-lab/temporal-shift-module](https://github.com/mit-han-lab/temporal-shift-module).
|
|
|
|
More details about the shift module(which is the core of TSM) could to [test_shift.py](./test_shift.py).
|
|
|
|
## Tutorial
|
|
|
|
+ An example could refer to [demo.sh](./demo.sh)
|
|
+ Requirements: Successfully installed `torch>=1.3.0, torchvision`
|
|
|
|
+ Step 1: Train/Download TSM-R50 checkpoints from [offical Github repo](https://github.com/mit-han-lab/temporal-shift-module) or [MMAction2](https://github.com/open-mmlab/mmaction2)
|
|
+ Supported settings: `num_segments`, `shift_div`, `num_classes`.
|
|
+ Fixed settings: `backbone`(ResNet50), `shift_place`(blockres), `temporal_pool`(False).
|
|
|
|
+ Step 2: Convert PyTorch checkpoints to TensorRT weights.
|
|
|
|
```shell
|
|
python gen_wts.py /path/to/pytorch.pth --out-filename /path/to/tensorrt.wts
|
|
```
|
|
|
|
+ Step 3: Test Python API.
|
|
+ Modify configs in `tsm_r50.py`.
|
|
+ Inference with `tsm_r50.py`.
|
|
|
|
```python
|
|
# Supported settings
|
|
BATCH_SIZE = 1
|
|
NUM_SEGMENTS = 8
|
|
INPUT_H = 224
|
|
INPUT_W = 224
|
|
OUTPUT_SIZE = 400
|
|
SHIFT_DIV = 8
|
|
```
|
|
|
|
```shell
|
|
usage: tsm_r50.py [-h] [--tensorrt-weights TENSORRT_WEIGHTS] [--input-video INPUT_VIDEO] [--save-engine-path SAVE_ENGINE_PATH] [--load-engine-path LOAD_ENGINE_PATH] [--test-mmaction2] [--mmaction2-config MMACTION2_CONFIG] [--mmaction2-checkpoint MMACTION2_CHECKPOINT] [--test-cpp] [--cpp-result-path CPP_RESULT_PATH]
|
|
|
|
optional arguments:
|
|
-h, --help show this help message and exit
|
|
--tensorrt-weights TENSORRT_WEIGHTS
|
|
Path to TensorRT weights, which is generated by gen_weights.py
|
|
--input-video INPUT_VIDEO
|
|
Path to local video file
|
|
--save-engine-path SAVE_ENGINE_PATH
|
|
Save engine to local file
|
|
--load-engine-path LOAD_ENGINE_PATH
|
|
Saved engine file path
|
|
--test-mmaction2 Compare TensorRT results with MMAction2 Results
|
|
--mmaction2-config MMACTION2_CONFIG
|
|
Path to MMAction2 config file
|
|
--mmaction2-checkpoint MMACTION2_CHECKPOINT
|
|
Path to MMAction2 checkpoint url or file path
|
|
--test-cpp Compare Python API results with C++ API results
|
|
--cpp-result-path CPP_RESULT_PATH
|
|
Path to C++ API results
|
|
```
|
|
|
|
+ Step 4: Test C++ API.
|
|
+ Mocify Configs in `tsm_r50.cpp`.
|
|
+ Build from source code: `mkdir build && cd build && cmake .. && make`
|
|
+ Generate Engine file: `./tsm_r50 -s`
|
|
+ Inference with genrated engine file and write predictions to local: `./tsm_r50 -d`
|
|
+ Compare results with Python API: `python tsm_r50.py --tensorrt-weights /path/to/tensorrt.weights --test-cpp --cpp-result-file /path/to/cpp-result.txt`
|
|
|
|
## TODO
|
|
|
|
+ [x] Python Shift module.
|
|
+ [x] Generate wts of official tsm and mmaction2 tsm.
|
|
+ [x] Python API Definition
|
|
+ [x] Test with mmaction2 demo
|
|
+ [x] Tutorial
|
|
+ [x] C++ API Definition
|