tensorrt_plugin.md 7.42 KB
Newer Older
RunningLeon's avatar
RunningLeon committed
1
## TensorRT Deployment
2

3
4
5
6
### <span style="color:red">DeprecationWarning</span>

TensorRT support will be deprecated in the future.
Welcome to use the unified model deployment toolbox MMDeploy: https://github.com/open-mmlab/mmdeploy
7

8
9
<!-- TOC -->

RunningLeon's avatar
RunningLeon committed
10
- [TensorRT Deployment](#tensorrt-deployment)
11
  - [<span style="color:red">DeprecationWarning</span>](#deprecationwarning)
12
13
14
15
16
17
18
19
20
21
22
23
24
25
  - [Introduction](#introduction)
  - [List of TensorRT plugins supported in MMCV](#list-of-tensorrt-plugins-supported-in-mmcv)
  - [How to build TensorRT plugins in MMCV](#how-to-build-tensorrt-plugins-in-mmcv)
    - [Prerequisite](#prerequisite)
    - [Build on Linux](#build-on-linux)
  - [Create TensorRT engine and run inference in python](#create-tensorrt-engine-and-run-inference-in-python)
  - [How to add a TensorRT plugin for custom op in MMCV](#how-to-add-a-tensorrt-plugin-for-custom-op-in-mmcv)
    - [Main procedures](#main-procedures)
    - [Reminders](#reminders)
  - [Known Issues](#known-issues)
  - [References](#references)

<!-- TOC -->

26
### Introduction
27
28
29
30

**NVIDIA TensorRT** is a software development kit(SDK) for high-performance inference of deep learning models. It includes a deep learning inference optimizer and runtime that delivers low latency and high-throughput for deep learning inference applications. Please check its [developer's website](https://developer.nvidia.com/tensorrt) for more information.
To ease the deployment of trained models with custom operators from `mmcv.ops` using TensorRT, a series of TensorRT plugins are included in MMCV.

31
### List of TensorRT plugins supported in MMCV
32

RunningLeon's avatar
RunningLeon committed
33
| ONNX Operator             | TensorRT Plugin                                                                 | MMCV Releases |
34
| :------------------------ | :------------------------------------------------------------------------------ | :-----------: |
RunningLeon's avatar
RunningLeon committed
35
36
37
38
39
40
41
| MMCVRoiAlign              | [MMCVRoiAlign](./tensorrt_custom_ops.md#mmcvroialign)                           |     1.2.6     |
| ScatterND                 | [ScatterND](./tensorrt_custom_ops.md#scatternd)                                 |     1.2.6     |
| NonMaxSuppression         | [NonMaxSuppression](./tensorrt_custom_ops.md#nonmaxsuppression)                 |     1.3.0     |
| MMCVDeformConv2d          | [MMCVDeformConv2d](./tensorrt_custom_ops.md#mmcvdeformconv2d)                   |     1.3.0     |
| grid_sampler              | [grid_sampler](./tensorrt_custom_ops.md#grid-sampler)                           |     1.3.1     |
| cummax                    | [cummax](./tensorrt_custom_ops.md#cummax)                                       |     1.3.5     |
| cummin                    | [cummin](./tensorrt_custom_ops.md#cummin)                                       |     1.3.5     |
42
| MMCVInstanceNormalization | [MMCVInstanceNormalization](./tensorrt_custom_ops.md#mmcvinstancenormalization) |     1.3.5     |
RunningLeon's avatar
RunningLeon committed
43
| MMCVModulatedDeformConv2d | [MMCVModulatedDeformConv2d](./tensorrt_custom_ops.md#mmcvmodulateddeformconv2d) |     1.3.8     |
SemyonBevzuk's avatar
SemyonBevzuk committed
44

45
46
47
48
Notes

- All plugins listed above are developed on TensorRT-7.2.1.6.Ubuntu-16.04.x86_64-gnu.cuda-10.2.cudnn8.0

49
### How to build TensorRT plugins in MMCV
50

51
#### Prerequisite
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

- Clone repository

```bash
git clone https://github.com/open-mmlab/mmcv.git
```

- Install TensorRT

Download the corresponding TensorRT build from [NVIDIA Developer Zone](https://developer.nvidia.com/nvidia-tensorrt-download).

For example, for Ubuntu 16.04 on x86-64 with cuda-10.2, the downloaded file is `TensorRT-7.2.1.6.Ubuntu-16.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz`.

Then, install as below:

```bash
cd ~/Downloads
tar -xvzf TensorRT-7.2.1.6.Ubuntu-16.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz
export TENSORRT_DIR=`pwd`/TensorRT-7.2.1.6
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TENSORRT_DIR/lib
```

Install python packages: tensorrt, graphsurgeon, onnx-graphsurgeon

```bash
pip install $TENSORRT_DIR/python/tensorrt-7.2.1.6-cp37-none-linux_x86_64.whl
pip install $TENSORRT_DIR/onnx_graphsurgeon/onnx_graphsurgeon-0.2.6-py2.py3-none-any.whl
pip install $TENSORRT_DIR/graphsurgeon/graphsurgeon-0.4.5-py2.py3-none-any.whl
```

82
For more detailed information of installing TensorRT using tar, please refer to [Nvidia' website](https://docs.nvidia.com/deeplearning/tensorrt/archives/tensorrt-721/install-guide/index.html#installing-tar).
83

84
85
86
87
- Install cuDNN

Install cuDNN 8 following [Nvidia' website](https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html#installlinux-tar).

88
#### Build on Linux
89
90

```bash
91
cd mmcv ## to MMCV root directory
92
93
94
MMCV_WITH_OPS=1 MMCV_WITH_TRT=1 pip install -e .
```

95
### Create TensorRT engine and run inference in python
96
97
98
99
100
101
102

Here is an example.

```python
import torch
import onnx

lizz's avatar
lizz committed
103
from mmcv.tensorrt import (TRTWrapper, onnx2trt, save_trt_engine,
104
105
106
107
108
109
110
111
                                   is_tensorrt_plugin_loaded)

assert is_tensorrt_plugin_loaded(), 'Requires to complie TensorRT plugins in mmcv'

onnx_file = 'sample.onnx'
trt_file = 'sample.trt'
onnx_model = onnx.load(onnx_file)

112
## Model input
113
inputs = torch.rand(1, 3, 224, 224).cuda()
114
## Model input shape info
115
116
117
118
119
120
opt_shape_dict = {
    'input': [list(inputs.shape),
              list(inputs.shape),
              list(inputs.shape)]
}

121
## Create TensorRT engine
122
123
124
125
126
127
max_workspace_size = 1 << 30
trt_engine = onnx2trt(
    onnx_model,
    opt_shape_dict,
    max_workspace_size=max_workspace_size)

128
## Save TensorRT engine
129
130
save_trt_engine(trt_engine, trt_file)

131
## Run inference with TensorRT
lizz's avatar
lizz committed
132
trt_model = TRTWrapper(trt_file, ['input'], ['output'])
133
134
135
136
137
138
139

with torch.no_grad():
    trt_outputs = trt_model({'input': inputs})
    output = trt_outputs['output']

```

140
### How to add a TensorRT plugin for custom op in MMCV
141

142
#### Main procedures
143
144
145
146
147
148
149
150
151
152
153
154

Below are the main steps:

1. Add c++ header file
2. Add c++ source file
3. Add cuda kernel file
4. Register plugin in `trt_plugin.cpp`
5. Add unit test in `tests/test_ops/test_tensorrt.py`

**Take RoIAlign plugin `roi_align` for example.**

1. Add header `trt_roi_align.hpp` to TensorRT include directory `mmcv/ops/csrc/tensorrt/`
155

156
2. Add source `trt_roi_align.cpp` to TensorRT source directory `mmcv/ops/csrc/tensorrt/plugins/`
157

158
3. Add cuda kernel `trt_roi_align_kernel.cu` to TensorRT source directory `mmcv/ops/csrc/tensorrt/plugins/`
159

160
161
4. Register `roi_align` plugin in [trt_plugin.cpp](https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/csrc/tensorrt/plugins/trt_plugin.cpp)

162
163
   ```c++
   #include "trt_plugin.hpp"
164

165
   #include "trt_roi_align.hpp"
166

167
   REGISTER_TENSORRT_PLUGIN(RoIAlignPluginDynamicCreator);
168

169
170
171
172
   extern "C" {
   bool initLibMMCVInferPlugins() { return true; }
   }  // extern "C"
   ```
173
174
175
176

5. Add unit test into `tests/test_ops/test_tensorrt.py`
   Check [here](https://github.com/open-mmlab/mmcv/blob/master/tests/test_ops/test_tensorrt.py) for examples.

177
#### Reminders
178

RunningLeon's avatar
RunningLeon committed
179
180
- *Please note that this feature is experimental and may change in the future. Strongly suggest users always try with the latest master branch.*

lizz's avatar
lizz committed
181
- Some of the [custom ops](https://mmcv.readthedocs.io/en/latest/ops.html) in `mmcv` have their cuda implementations, which could be referred.
182

183
### Known Issues
184
185
186

- None

187
### References
188
189
190
191
192
193

- [Developer guide of Nvidia TensorRT](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html)
- [TensorRT Open Source Software](https://github.com/NVIDIA/TensorRT)
- [onnx-tensorrt](https://github.com/onnx/onnx-tensorrt)
- [TensorRT python API](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/index.html)
- [TensorRT c++ plugin API](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_plugin.html)