Skip to content

Commit

Permalink
[Feature] Add tensorrt evaluation (open-mmlab#328)
Browse files Browse the repository at this point in the history 
* trt_test

* remove history

* remove history

* remove history

* use html table

* use html table

* use html table

* eval tensorrt

* merge

* merge into one table

* specify all the prerequistes

* render table

* fix lint

* minor style change
  • Loading branch information
@AllentDan
AllentDan authored May 31, 2021
1 parent 5d1639d commit 973c48b
Show file tree
Hide file tree
Showing 3 changed files with 189 additions and 65 deletions.
7 changes: 7 additions & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,3 +116,10 @@ work_dirs/

# Pytorch
*.pth

# onnx and tensorrt
*.onnx
*.trt

# local history
.history/**
163 changes: 99 additions & 64 deletions docs/tools_scripts.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,24 +99,74 @@ The table below lists the models that are guaranteed to be exportable to ONNX an
- If you meet any problem with the listed models above, please create an issue and it would be taken care of soon. For models not included in the list, please try to solve them by yourself.
- Because this feature is experimental and may change fast, please always try with the latest `mmcv` and `mmedit`.

### Evaluate ONNX Models with ONNXRuntime (experimental)
### Convert ONNX to TensorRT (experimental)

We also provide a script to convert [ONNX](https://github.com/onnx/onnx) model to [TensorRT](https://github.com/NVIDIA/TensorRT) format. Besides, we support comparing the output results between ONNX and TensorRT model.

We prepare a tool `tools/deploy_test.py` to evaluate ONNX models with ONNX Runtime backend.

#### Prerequisite
```bash
python tools/onnx2tensorrt.py
${CFG_PATH} \
${MODEL_TYPE} \
${IMAGE_PATH} \
${INPUT_ONNX} \
--trt-file ${OUT_TENSORRT} \
--max-shape INT INT INT INT \
--min-shape INT INT INT INT \
--workspace-size INT \
--fp16 \
--show \
--verify \
--verbose
```

Description of arguments:

- `config` : The path of a model config file.
- `model_type` :The model type of the config file, options: `inpainting`, `mattor`, `restorer`, `synthesizer`.
- `img_path` : The path to input image file.
- `onnx_file` : The path to input ONNX file.
- `--trt-file` : The path of output TensorRT model. If not specified, it will be set to `tmp.trt`.
- `--max-shape` : Maximum shape of model input.
- `--min-shape` : Minimum shape of model input.
- `--workspace-size`: Max workspace size in GiB. If not specified, it will be set to 1 GiB.
- `--fp16`: Determines whether to export TensorRT with fp16 mode. If not specified, it will be set to `False`.
- `--show`: Determines whether to show the output of ONNX and TensorRT. If not specified, it will be set to `False`.
- `--verify`: Determines whether to verify the correctness of an exported model. If not specified, it will be set to `False`.
- `--verbose`: Determines whether to verbose logging messages while creating TensorRT engine. If not specified, it will be set to `False`.

**Note**: This tool is still experimental. Some customized operators are not supported for now. We only support `restorer` for now. While generating ONNX file of SRCNN, replace 'bicubic' with 'bilinear' in SCRNN model [here](https://github.com/open-mmlab/mmediting/blob/764e6065e315b7d0033762038fcbf0bb1c570d4d/mmedit/models/backbones/sr_backbones/srcnn.py#L40). For TensorRT does not support bicubic interpolation by now and final performace will be weaken by about 4%.

#### List of supported models exportable to TensorRT

The table below lists the models that are guaranteed to be exportable to TensorRT engine and runnable in TensorRT.

- Install onnx and onnxruntime-gpu
| Model | Config | Dynamic Shape | Batch Inference | Note |
| :------: | :-----------------------------------------------------------------------------------------------------------------------------------------------------------------: | :-----------: | :-------------: | :---: |
| ESRGAN | [esrgan_x4c64b23g32_g1_400k_div2k.py](https://github.com/open-mmlab/mmediting/blob/master/configs/restorers/esrgan/esrgan_x4c64b23g32_g1_400k_div2k.py) | Y | Y | |
| ESRGAN | [esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py](https://github.com/open-mmlab/mmediting/blob/master/configs/restorers/esrgan/esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py) | Y | Y | |
| SRCNN | [srcnn_x4k915_g1_1000k_div2k.py](https://github.com/open-mmlab/mmediting/blob/master/configs/restorers/srcnn/srcnn_x4k915_g1_1000k_div2k.py) | Y | Y | 'bicubic' upsampling must be replaced with 'bilinear' |

```shell
pip install onnx onnxruntime-gpu
```
**Notes**:

- *All models above are tested with Pytorch==1.8.1, onnxruntime==1.7.0 and tensorrt==7.2.3.4*
- If you meet any problem with the listed models above, please create an issue and it would be taken care of soon. For models not included in the list, please try to solve them by yourself.
- Because this feature is experimental and may change fast, please always try with the latest `mmcv` and `mmedit`.

### Evaluate ONNX and TensorRT Models (experimental)

We provide methods to evaluate TensorRT and ONNX models in `tools/deploy_test.py`.

#### Prerequisite
To evaluate ONNX and TensorRT models, onnx, onnxruntime and TensorRT should be installed first. Install `mmcv-full` with ONNXRuntime custom ops and TensorRT plugins follow [ONNXRuntime in mmcv](https://mmcv.readthedocs.io/en/latest/onnxruntime_op.html) and [TensorRT plugin in mmcv](https://github.com/open-mmlab/mmcv/blob/master/docs/tensorrt_plugin.md).

#### Usage

```bash
python tools/deploy_test.py \
${CONFIG_FILE} \
${ONNX_FILE} \
${MODEL_PATH} \
${BACKEND} \
--out ${OUTPUT_FILE} \
--save-path ${SAVE_PATH} \
----cfg-options ${CFG_OPTIONS} \
Expand All @@ -125,7 +175,8 @@ python tools/deploy_test.py \
#### Description of all arguments

- `config`: The path of a model config file.
- `model`: The path of an ONNX model file.
- `model`: The path of a TensorRT or an ONNX model file.
- `backend`: The backend for testing, choose tensorrt or onnxruntime.
- `--out`: The path of output result file in pickle format.
- `--save-path`: The path to store images and if not given, it will not save image.
- `--cfg-options`: Override some settings in the used config file, the key-value pair in `xxx=yyy` format will be merged into config file.
Expand All @@ -140,6 +191,8 @@ python tools/deploy_test.py \
<th align="center">Metric</th>
<th align="center">PyTorch</th>
<th align="center">ONNX Runtime</th>
<th align="center">TensorRT FP32</th>
<th align="center">TensorRT FP16</th>
</tr>
<tr>
<td align="center" rowspan="6">ESRGAN</td>
Expand All @@ -150,33 +203,45 @@ python tools/deploy_test.py \
<td align="center">PSNR</td>
<td align="center">28.2700</td>
<td align="center">28.2619</td>
<td align="center">28.2619</td>
<td align="center">28.2616</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.7778</td>
<td align="center">0.7784</td>
<td align="center">0.7784</td>
<td align="center">0.7783</td>
</tr>
<tr>
<td align="center" rowspan="2">Set14</td>
<td align="center">PSNR</td>
<td align="center">24.6328</td>
<td align="center">24.6290</td>
<td align="center">24.6290</td>
<td align="center">24.6274</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.6491</td>
<td align="center">0.6494</td>
<td align="center">0.6494</td>
<td align="center">0.6494</td>
</tr>
<tr>
<td align="center" rowspan="2">DIV2K</td>
<td align="center">PSNR</td>
<td align="center">26.6531</td>
<td align="center">26.6532</td>
<td align="center">26.6532</td>
<td align="center">26.6532</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.7340</td>
<td align="center">0.7340</td>
<td align="center">0.7340</td>
<td align="center">0.7340</td>
</tr>
<tr>
<td align="center" rowspan="6">ESRGAN</td>
Expand All @@ -187,33 +252,45 @@ python tools/deploy_test.py \
<td align="center">PSNR</td>
<td align="center">30.6428</td>
<td align="center">30.6307</td>
<td align="center">30.6307</td>
<td align="center">30.6305</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.8559</td>
<td align="center">0.8565</td>
<td align="center">0.8565</td>
<td align="center">0.8566</td>
</tr>
<tr>
<td align="center" rowspan="2">Set14</td>
<td align="center">PSNR</td>
<td align="center">27.0543</td>
<td align="center">27.0422</td>
<td align="center">27.0422</td>
<td align="center">27.0411</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.7447</td>
<td align="center">0.7450</td>
<td align="center">0.7450</td>
<td align="center">0.7449</td>
</tr>
<tr>
<td align="center" rowspan="2">DIV2K</td>
<td align="center">PSNR</td>
<td align="center">29.3354</td>
<td align="center">29.3354</td>
<td align="center">29.3354</td>
<td align="center">29.3339</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.8263</td>
<td align="center">0.8263</td>
<td align="center">0.8263</td>
<td align="center">0.8263</td>
</tr>
<tr>
<td align="center" rowspan="6">SRCNN</td>
Expand All @@ -224,91 +301,49 @@ python tools/deploy_test.py \
<td align="center">PSNR</td>
<td align="center">28.4316</td>
<td align="center">28.4120</td>
<td align="center">27.2144</td>
<td align="center">27.2127</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.8099</td>
<td align="center">0.8106</td>
<td align="center">0.7782</td>
<td align="center">0.7781</td>
</tr>
<tr>
<td align="center" rowspan="2">Set14</td>
<td align="center">PSNR</td>
<td align="center">25.6486</td>
<td align="center">25.6367</td>
<td align="center">24.8613</td>
<td align="center">24.8599</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.7014</td>
<td align="center">0.7015</td>
<td align="center">0.6674</td>
<td align="center">0.6673</td>
</tr>
<tr>
<td align="center" rowspan="2">DIV2K</td>
<td align="center">PSNR</td>
<td align="center">27.7460</td>
<td align="center">27.7460</td>
<td align="center">26.9891</td>
<td align="center">26.9862</td>
</tr>
<tr>
<td align="center">SSIM</td>
<td align="center">0.7854</td>
<td align="center">0.78543</td>
<td align="center">0.7605</td>
<td align="center">0.7604</td>
</tr>
</table>

**Notes**:

- All ONNX models are evaluated with dynamic shape on the datasets and images are preprocessed according to the original config file.
- All ONNX and TensorRT models are evaluated with dynamic shape on the datasets and images are preprocessed according to the original config file.
- This tool is still experimental, and we only support `restorer` for now.


### Convert ONNX to TensorRT (experimental)
We also provide a script to convert [ONNX](https://github.com/onnx/onnx) model to [TensorRT](https://github.com/NVIDIA/TensorRT) format. Besides, we support comparing the output results between ONNX and TensorRT model.


```bash
python tools/onnx2tensorrt.py
${CFG_PATH} \
${MODEL_TYPE} \
${IMAGE_PATH} \
${INPUT_ONNX} \
--trt-file ${OUT_TENSORRT} \
--max-shape INT INT INT INT \
--min-shape INT INT INT INT \
--workspace-size INT \
--fp16 \
--show \
--verify \
--verbose
```

Description of arguments:

- `config` : The path of a model config file.
- `model_type` :The model type of the config file, options: `inpainting`, `mattor`, `restorer`, `synthesizer`.
- `img_path` : The path to input image file.
- `onnx_file` : The path to input ONNX file.
- `--trt-file` : The path of output TensorRT model. If not specified, it will be set to `tmp.trt`.
- `--max-shape` : Maximum shape of model input.
- `--min-shape` : Minimum shape of model input.
- `--workspace-size`: Max workspace size in GiB. If not specified, it will be set to 1 GiB.
- `--fp16`: Determines whether to export TensorRT with fp16 mode. If not specified, it will be set to `False`.
- `--show`: Determines whether to show the output of ONNX and TensorRT. If not specified, it will be set to `False`.
- `--verify`: Determines whether to verify the correctness of an exported model. If not specified, it will be set to `False`.
- `--verbose`: Determines whether to verbose logging messages while creating TensorRT engine. If not specified, it will be set to `False`.

**Note**: This tool is still experimental. Some customized operators are not supported for now. We only support `restorer` for now. While generating ONNX file of SRCNN, replace 'bicubic' with 'bilinear' in SCRNN model [here](https://github.com/open-mmlab/mmediting/blob/764e6065e315b7d0033762038fcbf0bb1c570d4d/mmedit/models/backbones/sr_backbones/srcnn.py#L40). For TensorRT does not support bicubic interpolation by now and final performace will be weaken by about 4%.

#### List of supported models exportable to TensorRT

The table below lists the models that are guaranteed to be exportable to TensorRT engine and runnable in TensorRT.

| Model | Config | Dynamic Shape | Batch Inference | Note |
| :------: | :-----------------------------------------------------------------------------------------------------------------------------------------------------------------: | :-----------: | :-------------: | :---: |
| ESRGAN | [esrgan_x4c64b23g32_g1_400k_div2k.py](https://github.com/open-mmlab/mmediting/blob/master/configs/restorers/esrgan/esrgan_x4c64b23g32_g1_400k_div2k.py) | Y | Y | |
| ESRGAN | [esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py](https://github.com/open-mmlab/mmediting/blob/master/configs/restorers/esrgan/esrgan_psnr_x4c64b23g32_g1_1000k_div2k.py) | Y | Y | |
| SRCNN | [srcnn_x4k915_g1_1000k_div2k.py](https://github.com/open-mmlab/mmediting/blob/master/configs/restorers/srcnn/srcnn_x4k915_g1_1000k_div2k.py) | Y | Y | 'bicubic' upsampling must be replaced with 'bilinear' |

**Notes**:

- *All models above are tested with Pytorch==1.8.1, onnxruntime==1.7.0 and tensorrt==7.2.3.4*
- If you meet any problem with the listed models above, please create an issue and it would be taken care of soon. For models not included in the list, please try to solve them by yourself.
- Because this feature is experimental and may change fast, please always try with the latest `mmcv` and `mmedit`.
Loading

0 comments on commit 973c48b

Please sign in to comment.