!3085 ONNX:simple_baselines

Merge pull request !3085 from 谈宇隆/onnx_simplebase

research/cv/simple_baselines/README_CN.md

@@ -14,6 +14,7 @@
     - [脚本参数](#脚本参数)
     - [训练过程](#训练过程)
     - [评估过程](#评估过程)
+        - [onnx推理](#onnx推理)
 - [模型描述](#模型描述)
     - [性能](#性能)
         - [评估性能](#评估性能)
@@ -103,6 +104,7 @@ simple_baselines的总体网络架构如下：
     ├── run_distribute_train.sh            # 启动Ascend分布式训练（8卡）
     ├── run_eval.sh                        # 启动Ascend评估
     ├── run_standalone_train.sh            # 启动Ascend单机训练（单卡）
+    ├── run_onnx_eval.sh                   # 启动onnx推理
   ├── src
     ├── utils
         ├── coco.py                        # COCO数据集评估结果
@@ -114,6 +116,7 @@ simple_baselines的总体网络架构如下：
     ├── network_with_loss.py               # 损失函数定义
     └── pose_resnet.py                     # 主干网络定义
   ├── eval.py                              # 评估网络
+  ├── eval_onnx.py                         # onnx推理
   └── train.py                             # 训练网络
 ```
 
@@ -231,6 +234,40 @@ coco eval results saved to /cache/train_output/multi_train_poseresnet_v5_2-140_2
 AP: 0.704
 ```
 
+### onnx推理
+
+在推理之前需要先导出模型，onnx模型可在任意环境导出,通过device_target参数可进行设置。
+
+```bash
+python export.py --device_target CPU --ckpt_url /path/to/exported.ckpt --file_format ONNX
+```
+
+推理数据集目录为下
+
+```text
+└──coco
+  ├── val2017
+    ├── 000000000139.jpg
+    ├── ...
+    └── 000000000285.jpg
+  ├── annotations
+    └── person_keypoints_val2017.json
+```
+
+推理前需要更改src/config.py文件。更改config.DATASET.ROOT为当前数据集coco目录，例如/home/dataset/coco/，更改config.DATASET.TEST_JSON为annotations/person_keypoints_val2017.json
+
+```bash
+# onnx模型推理
+bash run_onnx_eval.sh  CKPT_FILES DEVICE_TYPE
+```
+
+推理结果保存在output.eval_onnx.log文件下
+simple_baselines网络使用COCO数据集推理得到的结果如下：
+
+```text
+AP:0.72296
+```
+
 ## 推理过程
 
 ### [导出mindir]
@@ -261,7 +298,7 @@ python export.py
 # You will see simple_pose.mindir under {Output file path}.
 ```
 
-`FILE_FORMAT` 变量应该在["AIR", "MINDIR"]中选择
+`FILE_FORMAT` 变量应该在["AIR", "MINDIR", "ONNX"]中选择
 
 ### 310推理
 

research/cv/simple_baselines/eval_onnx.py

+# 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.
+# ============================================================================
+"""Run evaluation for a model exported to ONNX"""
+import time
+import argparse
+import os
+import numpy as np
+import onnxruntime as ort
+from tqdm import tqdm
+from src.config import config
+from src.dataset import keypoint_dataset
+from src.dataset import flip_pairs
+from src.utils.coco import evaluate
+from src.utils.transforms import flip_back
+from src.utils.inference import get_final_preds
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Evaluate onnx')
+    parser.add_argument('--ckpt_path', type=str, default="E:/simple_baselines/simple_baselines.onnx", help='onnx ckpt')
+    parser.add_argument('--target_device', type=str, default="GPU", help='select[GPU,CPU]')
+    args = parser.parse_args()
+    return args
+
+
+def create_session(checkpoint_path, target_device):
+    if target_device == 'GPU':
+        providers = ['CUDAExecutionProvider']
+    elif target_device == 'CPU':
+        providers = ['CPUExecutionProvider']
+    else:
+        raise ValueError(
+            f'Unsupported target device {target_device}, '
+            f'Expected one of: "CPU", "GPU"'
+        )
+    session = ort.InferenceSession(checkpoint_path, providers=providers)
+    input_name = session.get_inputs()[0].name
+    return session, input_name
+
+def validate(cfg, val_dataset, output_dir, ann_path, ckpt_path, target_device):
+    num_samples = val_dataset.get_dataset_size() * cfg.TEST.BATCH_SIZE
+    all_preds = np.zeros((num_samples, cfg.MODEL.NUM_JOINTS, 3),
+                         dtype=np.float32)
+    all_boxes = np.zeros((num_samples, 2))
+    image_id = []
+    idx = 0
+
+    start = time.time()
+    session, input_name = create_session(ckpt_path, target_device)
+    for item in tqdm(val_dataset.create_dict_iterator()):
+        inputs = item['image'].asnumpy()
+        output = session.run(None, {input_name: inputs})[0]
+        if cfg.TEST.FLIP_TEST:
+            inputs_flipped = inputs[:, :, :, ::-1]
+            output_flipped = session.run(None, {input_name: inputs_flipped})[0]
+            output_flipped = flip_back(output_flipped, flip_pairs)
+            if cfg.TEST.SHIFT_HEATMAP:
+                output_flipped[:, :, :, 1:] = \
+                    output_flipped.copy()[:, :, :, 0:-1]
+            output = (output + output_flipped) * 0.5
+
+        c = item['center'].asnumpy()
+        s = item['scale'].asnumpy()
+        score = item['score'].asnumpy()
+        file_id = list(item['id'].asnumpy())
+        preds, maxvals = get_final_preds(cfg, output.copy(), c, s)
+        num_images, _ = preds.shape[:2]
+        all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
+        all_preds[idx:idx + num_images, :, 2:3] = maxvals
+        all_boxes[idx:idx + num_images, 0] = np.prod(s * 200, 1)
+        all_boxes[idx:idx + num_images, 1] = score
+        image_id.extend(file_id)
+        idx += num_images
+        if idx % 1024 == 0:
+            print('{} samples validated in {} seconds'.format(idx, time.time() - start))
+            start = time.time()
+
+
+    print(all_preds[:idx].shape, all_boxes[:idx].shape, len(image_id))
+    _, perf_indicator = evaluate(cfg, all_preds[:idx], output_dir, all_boxes[:idx], image_id, ann_path)
+    print("AP:", perf_indicator)
+    return perf_indicator
+
+def main():
+    args = parse_args()
+    valid_dataset, _ = keypoint_dataset(
+        config=config,
+        train_mode=False,
+        num_parallel_workers=config.TEST.NUM_PARALLEL_WORKERS,
+    )
+    ckpt_name = args.ckpt_path.split('/')
+    ckpt_name = ckpt_name[len(ckpt_name) - 1]
+    ckpt_name = ckpt_name.split('.')[0]
+
+    output_dir = os.path.join(config.DATASET.ROOT, 'result/')
+    ann_path = config.DATASET.ROOT
+    output_dir = output_dir + ckpt_name
+    ann_path = ann_path + config.DATASET.TEST_JSON
+    print("ann_path is :", ann_path)
+    validate(config, valid_dataset, output_dir, ann_path, args.ckpt_path, args.target_device)
+
+if __name__ == '__main__':
+    main()

research/cv/simple_baselines/export.py

@@ -31,7 +31,7 @@ parser.add_argument("--device_id", type=int, default=0, help="Device id")
 parser.add_argument("--ckpt_url", default="/home/dataset/coco/multi_train_poseresnet_commit_0-140_292.ckpt",
                     help="Checkpoint file path.")
 parser.add_argument("--file_name", type=str, default="simple_baselines", help="output file name.")
-parser.add_argument('--file_format', type=str, choices=["MINDIR"], default='MINDIR', help='file format')
+parser.add_argument('--file_format', type=str, choices=["MINDIR", "ONNX"], default='MINDIR', help='file format')
 args = parser.parse_args()
 
 if __name__ == '__main__':

research/cv/simple_baselines/scripts/run_onnx_eval.sh

+#!/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.
+# ============================================================================
+
+echo "=============================================================================================================="
+echo "Please run the script as: "
+echo "bash scripts/run_onnx_eval.sh CKPT_FILES  DEVICE_TYPE "
+echo "for example: bash scripts/run_onnx_eval.sh /path/a.onnx GPU "
+echo "=============================================================================================================="
+
+if [ $# -lt 2 ]
+then
+    echo "Usage: bash scripts/run_onnx_eval.sh [CKPT_FILES] [DEVICE_TYPE]"
+    exit 1
+fi
+
+CKPT_FILES=$1
+DEVICE_TYPE=$2
+
+python eval_onnx.py --ckpt_path=$CKPT_FILES \
+    --target_device=$DEVICE_TYPE > output.eval_onnx.log 2>&1 &