From d46c7d9eba89a0cecbf87824b3e7d3dc80db08c6 Mon Sep 17 00:00:00 2001
From: zqh962903415 <962903415@qq.com>
Date: Mon, 25 Jul 2022 15:48:38 +0800
Subject: [PATCH] ONNX converter: CNN+CTC support
add final newline
add final newline
add onnx infer
---
research/cv/STGAN/README.md | 34 ++++++++-
research/cv/STGAN/infer_stgan_onnx.py | 77 +++++++++++++++++++++
research/cv/STGAN/requirements.txt | 3 +-
research/cv/STGAN/scripts/run_infer_onnx.sh | 31 +++++++++
research/cv/STGAN/src/utils/args.py | 2 +
5 files changed, 145 insertions(+), 2 deletions(-)
create mode 100644 research/cv/STGAN/infer_stgan_onnx.py
create mode 100644 research/cv/STGAN/scripts/run_infer_onnx.sh
diff --git a/research/cv/STGAN/README.md b/research/cv/STGAN/README.md
index e3a648c7e..fee020019 100644
--- a/research/cv/STGAN/README.md
+++ b/research/cv/STGAN/README.md
@@ -104,6 +104,8 @@ After installing MindSpore via the official website, you can start training and
鈹� 鈹溾攢鈹€run_standalone_train_gpu.sh // train in GPU
鈹� 鈹溾攢鈹€run_eval_gpu.sh // evaluate in GPU
鈹� 鈹溾攢鈹€run_distribute_train_gpu.sh // distributed train in GPU
+ 鈹� 鈹溾攢鈹€run_infer_onnx.sh // infer onnx in GPU
+ 鈹� 鈹溾攢鈹€run_infer_310.sh // infer in ascend
鈹溾攢鈹€ src
鈹溾攢鈹€ dataset
鈹溾攢鈹€ datasets.py // creating dataset
@@ -120,6 +122,9 @@ After installing MindSpore via the official website, you can start training and
鈹溾攢鈹€ train.py // training script
鈹溾攢鈹€ eval.py // evaluation script
鈹溾攢鈹€ export.py // model-export script
+ 鈹溾攢鈹€ infer_stgan_onnx.py // onnx model infer script
+ 鈹溾攢鈹€ postprocess.py // post process for 310 inference
+ 鈹溾攢鈹€ preprocess.py // pre process for 310 inference
```
### [Script Parameters](#contents)
@@ -132,6 +137,8 @@ Major parameters in train.py and utils/args.py as follows:
--batch_size: Training batch size.
--image_size: Image size used as input to the model.
--device_target: Device where the code will be implemented. Optional value is "Ascend" or "GPU".
+--onnx_path: Path of onnx file.
+--onnx_output: Path of onnx infer output.
```
### [Training Process](#contents)
@@ -202,7 +209,10 @@ Before running the command below, please check the checkpoint path used for eval
python export.py --ckpt_path [CHECKPOINT_PATH] --platform [PLATFORM] --file_format[EXPORT_FORMAT]
```
-If you want to infer the network on Ascend 310, `EXPORT_FORMAT` should be "MINDIR"
+- The ckpt_path parameter is required.
+- `PLATFORM` should be in ["Ascend", "GPU", "CPU"]
+- `EXPORT_FORMAT` should be in ["AIR", "ONNX", "MINDIR"].
+- If you want to infer the network on Ascend 310, `EXPORT_FORMAT` should be "MINDIR"
## Inference Process
@@ -231,6 +241,23 @@ bash run_infer_310.sh [GEN_MINDIR_PATH] [DATA_PATH][NEED_PREPROCESS] [DEVICE_ID]
- `NEED_PREPROCESS` means weather need preprocess or not, it's value is 'y' or 'n'. This step will process the image and label into .bin file and put them in the `process_Data` folder.
- `DEVICE_ID` is optional, it can be set by environment variable device_id, otherwise the value is zero.
+### ONNX Infer
+
+Before inferring, the onnx file should be exported by the `export.py` script
+
+```shell
+# ONNX inference
+python infer_stgan_onnx.py --dataroot ./dataset --onnx_path ./STGAN_model.onnx --experiment_name 128 --platform GPU
+# or run the script
+bash ./scripts/run_infer_onnx.sh [DATA_PATH] [EXPERIMENT_NAME] [DEVICE_ID] [ONNX_PATH]
+eg: bash ./scripts/run_infer_onnx.sh ./dataset 128 0 ./STGAN_model.onnx
+```
+
+- `DATA_PATH` path of dataset
+- `EXPERIMENT_NAME` default:128
+- `DEVICE_ID` default:0
+- `ONNX_PATH` onnx file path
+
### Result
Inference result is saved in `result_Files/` in current path, Inference time result is saved in `time_Result/`. The edited picture is saved as xxx.jpg format, such as `183800.jpg`.
@@ -240,6 +267,11 @@ Inference result is saved in `result_Files/` in current path, Inference time res
NN inference cost average time: 9.98606 ms of infer_count 10
```
+- ONNX result
+
+Onnx infer outputs are saved in `infer_onnx_outputs/`
+Onnx infer results is saved in `infer_onnx_log.log`
+
## Model Description
### Performance
diff --git a/research/cv/STGAN/infer_stgan_onnx.py b/research/cv/STGAN/infer_stgan_onnx.py
new file mode 100644
index 000000000..528fb1d54
--- /dev/null
+++ b/research/cv/STGAN/infer_stgan_onnx.py
@@ -0,0 +1,77 @@
+# Copyright 2022 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+""" infer STGAN ONNX"""
+import os
+import onnxruntime
+import numpy as np
+import tqdm
+import cv2
+
+from mindspore.common import set_seed
+
+from src.utils import get_args
+from src.dataset import CelebADataLoader
+
+set_seed(1)
+
+
+def get_input_name(onnx_session):
+ """
+ input_name = onnx_session.get_inputs()[0].name
+ :param onnx_session:
+ :return:
+ """
+ input_name = []
+ for node in onnx_session.get_inputs():
+ input_name.append(node.name)
+ return input_name
+
+
+def run_eval():
+ """ eval onnx model function """
+ args = get_args("test")
+ print('\n\n=============== start eval onnx model ===============\n\n')
+ data_loader = CelebADataLoader(args.dataroot,
+ mode=args.phase,
+ selected_attrs=args.attrs,
+ batch_size=1,
+ image_size=args.image_size)
+ iter_per_epoch = len(data_loader)
+ args.dataset_size = iter_per_epoch
+ ## onnx model export
+ onnx_path = args.onnx_path
+ session = onnxruntime.InferenceSession(onnx_path)
+
+ for _ in tqdm.trange(iter_per_epoch, desc='Eval onnx model Loop'):
+
+ data = next(data_loader.test_loader)
+ input_image = data['image'].asnumpy()
+ input_label = data['label'].asnumpy()
+ filename = data_loader.test_set.get_current_filename()
+ fake_image = session.run(None, {get_input_name(session)[0]: input_image,
+ get_input_name(session)[1]: input_label})
+ images = fake_image[0]
+ final_imgs = np.transpose(images, (0, 2, 3, 1))
+ final_imgs = np.squeeze(final_imgs)
+ out_path = args.onnx_output
+ if not os.path.exists(out_path):
+ os.mkdir(out_path)
+ cv2.imwrite(os.path.join(out_path, str(filename)), cv2.cvtColor(final_imgs * 255, cv2.COLOR_BGR2RGB))
+
+
+ print('\n\n=============== finish eval onnx model ===============\n\n')
+
+if __name__ == '__main__':
+ run_eval()
diff --git a/research/cv/STGAN/requirements.txt b/research/cv/STGAN/requirements.txt
index 8a2d9729d..deb64cb70 100644
--- a/research/cv/STGAN/requirements.txt
+++ b/research/cv/STGAN/requirements.txt
@@ -1,2 +1,3 @@
numpy
-tqdm
\ No newline at end of file
+tqdm
+onnxruntime
diff --git a/research/cv/STGAN/scripts/run_infer_onnx.sh b/research/cv/STGAN/scripts/run_infer_onnx.sh
new file mode 100644
index 000000000..49dea17de
--- /dev/null
+++ b/research/cv/STGAN/scripts/run_infer_onnx.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# Copyright 2022 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+if [ $# != 4 ]
+then
+ echo "Usage: bash run_infer_onnx.sh [DATA_PATH] [EXPERIMENT_NAME] [DEVICE_ID] [ONNX_PATH]"
+exit 1
+fi
+
+
+export DATA_PATH=$1
+export EXPERIMENT_NAME=$2
+export DEVICE_ID=$3
+export ONNX_PATH=$4
+
+python infer_stgan_onnx.py --dataroot=$DATA_PATH --experiment_name=$EXPERIMENT_NAME \
+ --device_id=$DEVICE_ID --onnx_path=$ONNX_PATH --platform="GPU" > infer_onnx_log 2>&1 &
\ No newline at end of file
diff --git a/research/cv/STGAN/src/utils/args.py b/research/cv/STGAN/src/utils/args.py
index f1bd1d9c3..6628341b6 100644
--- a/research/cv/STGAN/src/utils/args.py
+++ b/research/cv/STGAN/src/utils/args.py
@@ -259,6 +259,8 @@ def get_args(phase):
help='file format')
parser.add_argument('--file_name', type=str, default='STGAN', help='output file name prefix.')
parser.add_argument('--ckpt_path', default=None, help='path of checkpoint file.')
+ parser.add_argument('--onnx_path', default=None, help='path of onnx file.')
+ parser.add_argument('--onnx_output', default='./infer_onnx_outputs', help='path of onnx infer output.')
args = parser.parse_args()
if phase == 'test':
--
GitLab