Support TSM-R50 C++ API (wang-xinyu#500)

* C++ API for TSM-R50

* fix bugs

Author: irvingzhang0512

tsm/CMakeLists.txt

@@ -0,0 +1,25 @@
+cmake_minimum_required(VERSION 2.6)
+
+project(TSM)
+
+add_definitions(-std=c++11)
+
+option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
+set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_BUILD_TYPE Debug)
+
+include_directories(${PROJECT_SOURCE_DIR}/include)
+# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
+# cuda
+include_directories(/usr/local/cuda/include)
+link_directories(/usr/local/cuda/lib64)
+
+# tensorrt
+include_directories(/home/ubuntu/TensorRT/include/)
+link_directories(/home/ubuntu/TensorRT/lib/)
+
+add_executable(tsm_r50 ${PROJECT_SOURCE_DIR}/tsm_r50.cpp)
+target_link_libraries(tsm_r50 nvinfer)
+target_link_libraries(tsm_r50 cudart)
+
+add_definitions(-O2 -pthread)

tsm/README.md

@@ -23,9 +23,12 @@ More details about the shift module(which is the core of TSM) could to [test_shi
 python gen_wts.py /path/to/pytorch.pth --out-filename /path/to/tensorrt.wts
 ```
 
-+ Step 3: Modify configs in `tsm_r50.py`.
++ 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
@@ -34,10 +37,8 @@ OUTPUT_SIZE = 400
 SHIFT_DIV = 8
 ```
 
-+ Step 4: Inference with `tsm_r50.py`.
-
 ```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]
+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
@@ -54,13 +55,23 @@ optional arguments:
                         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
-+ [ ] C++ API Definition
++ [x] C++ API Definition

tsm/demo.sh

@@ -5,16 +5,16 @@ wget https://download.openmmlab.com/mmaction/recognition/tsm/tsm_r50_1x1x8_50e_k
 # Step 2: Convert pytorch checkpoints to TensorRT weights
 python gen_wts.py tsm_r50_1x1x8_50e_kinetics400_rgb_20200607-af7fb746.pth --out-filename ./tsm_r50_kinetics400_mmaction2.wts
 
-# Step 3: Skip this step since we use default settings.
-
-# Step 4: Inference
-# 1) Save local engine file to `./tsm_r50_kinetics400_mmaction2.trt`.
+# Step 3: Test Python API.
+# 3.1 Skip this step since we use default settings.
+# 3.2 Inference
+# 3.2.1 Save local engine file to `./tsm_r50_kinetics400_mmaction2.trt`.
 python tsm_r50.py \
     --tensorrt-weights ./tsm_r50_kinetics400_mmaction2.wts \
     --save-engine-path ./tsm_r50_kinetics400_mmaction2.trt
 
-# 2) Predict the recognition result using a single video `demo.mp4`.
-#    Should print `Result class id 6`, aka `arm wrestling`
+# 3.2.2 Predict the recognition result using a single video `demo.mp4`.
+#       Should print `Result class id 6`, aka `arm wrestling`
 # Download demo video
 wget https://raw.githubusercontent.com/open-mmlab/mmaction2/master/demo/demo.mp4
 # # use *.wts as input
@@ -24,8 +24,8 @@ wget https://raw.githubusercontent.com/open-mmlab/mmaction2/master/demo/demo.mp4
 python tsm_r50.py --load-engine-path ./tsm_r50_kinetics400_mmaction2.trt \
     --input-video ./demo.mp4
 
-# 3) Optional: Compare inference result with MMAction2 TSM-R50 model
-#    Have to install MMAction2 First, please refer to https://github.com/open-mmlab/mmaction2/blob/master/docs/install.md
+# 3.2.3 Optional: Compare inference result with MMAction2 TSM-R50 model
+#       Have to install MMAction2 First, please refer to https://github.com/open-mmlab/mmaction2/blob/master/docs/install.md
 # pip3 install pytest-runner
 # pip3 install mmcv
 # pip3 install mmaction2
@@ -41,3 +41,15 @@ python tsm_r50.py --load-engine-path ./tsm_r50_kinetics400_mmaction2.trt \
 #     --test-mmaction2 \
 #     --mmaction2-config mmaction2_tsm_r50_config.py \
 #     --mmaction2-checkpoint tsm_r50_1x1x8_50e_kinetics400_rgb_20200607-af7fb746.pth
+
+# Step 4: Test Python API.
+# 4.1 Skip this step since we use default settings.
+# 4.2 Build CPP
+mkdir build && cd build && cmake .. && make
+# 4.3 Generate Engine file
+./tsm_r50 -s
+# 4.4 Get Predictions
+./tsm_r50 -d
+# 4.5 Compare C++ Results with Python Results
+cd ..
+python tsm_r50.py --test-cpp --tensorrt-weights ./tsm_r50_kinetics400_mmaction2.wts