[Feature] add onnxruntime test tool

docs/tools_scripts.md

@@ -79,3 +79,39 @@ Description of arguments:
 - `--dynamic-export`: Determines whether to export ONNX model with dynamic input and output shapes. If not specified, it will be set to `False`.
 
 **Note**: This tool is still experimental. Some customized operators are not supported for now. And we only support `mattor` and `restorer` for now.
+
+### Evaluate ONNX model with ONNXRuntime
+
+We provide `tools/ort_test.py` to evaluate ONNX model with ONNXRuntime backend.
+
+Install onnxruntime-gpu before you start the test.
+
+  ```bash
+  pip install onnxruntime-gpu
+  ```
+
+Usage:
+
+  ```bash
+  python ${MMEDIT_PATH}/tools/ort_test.py \
+        ${CFG_PATH} \
+        ${ONNX_PATH} \
+        --out ${OUT_PATH} \
+        --save-path ${SHOW_DIR}
+  ```
+
+Description of arguments:
+
+- `config` : The path of a model config file.
+- `model` : Input model file.
+- `--out` : Output result pickle file.
+- `--save-path` : Path to store images.
+
+**Note**: Only support `mattor` and `restorer` for now.
+
+Results and Models
+
+| Model  |               Config                | Dataset |   Metric    |    PyTorch    |  ONNXRuntime  |
+| :----: | :---------------------------------: | :-----: | :---------: | :-----------: | :-----------: |
+| SRCNN  |   srcnn_x4k915_g1_1000k_div2k.py    |  Set5   | PSNR / SSIM | 28.43 / 0.809 | 28.41 / 0.810 |
+| ESRGAN | esrgan_x4c64b23g32_g1_400k_div2k.py |  Set5   | PSNR / SSIM | 28.27 / 0.777 | 28.26 / 0.778 |

tools/ort_test.py

@@ -1,19 +1,19 @@
 import argparse
-import mmcv
-import onnxruntime as ort
 import os
 import os.path as osp
-import torch
-from torch import nn
 import warnings
+
+import mmcv
+import numpy as np
+import onnxruntime as ort
+import torch
 from mmcv.parallel import MMDataParallel
 from mmcv.runner import get_dist_info
-import numpy as np
+from torch import nn
 
 from mmedit.apis import single_gpu_test
 from mmedit.datasets import build_dataloader, build_dataset
-from mmedit.models import BaseMattor, BasicRestorer
-from mmedit.models import build_model
+from mmedit.models import BaseMattor, BasicRestorer, build_model
 
 
 def inference_with_session(sess, io_binding, output_names, input_tensor):
@@ -46,7 +46,7 @@ def forward(self,
                 merged,
                 trimap,
                 meta,
-                test_mode=False, 
+                test_mode=False,
                 save_image=False,
                 save_path=None,
                 iteration=None):
@@ -115,8 +115,10 @@ def __init__(self, onnx_file, cfg, device_id):
         options = [{}]
         is_cuda_available = ort.get_device() == 'GPU'
         if is_cuda_available:
-            providers.insert(0, 'CUDAExecutionProvider')
-            options.insert(0, {'device_id': device_id})
+            # providers.insert(0, 'CUDAExecutionProvider')
+            # options.insert(0, {'device_id': device_id})
+            providers.append('CUDAExecutionProvider')
+            options.append({'device_id': device_id})
 
         sess.set_providers(providers, options)
 
@@ -145,10 +147,7 @@ def parse_args():
     parser.add_argument('model', help='Input model file')
     parser.add_argument('--out', help='output result pickle file')
     parser.add_argument(
-        '--save-path',
-        default=None,
-        type=str,
-        help='path to store images and if not given, will not save image')
+        '--save-path', default=None, type=str, help='path to store images')
     parser.add_argument('--local_rank', type=int, default=0)
     args = parser.parse_args()
     if 'LOCAL_RANK' not in os.environ: