diff --git a/Classification/cnns/README.md b/Classification/cnns/README.md index da4a99efef03d95fa750abbb1012e01ed47aa968..164b08ed08053e6c074d576cabc9412cac597aaa 100644 --- a/Classification/cnns/README.md +++ b/Classification/cnns/README.md @@ -516,3 +516,33 @@ python3 preprocess_imagenet_validation_data.py ../data/imagenet/validation 至此，已经完成了全部的数据预处理，您可以直接跳转至**转换训练集**和**转换验证集**部分，轻松完成ImageNet-2012数据集到OFRecord的转换过程了。 +### OneFlow 模型转 ONNX 模型 + +ONNX (Open Neural Network Exchange) 是一种较为广泛使用的神经网络中间格式，通过 ONNX 格式，OneFlow 模型可以被许多部署框架（如 OpenVINO、ONNX Runtime 和移动端的 ncnn、tnn、TEngine 等）所使用。这一节介绍如何将训练好的 resnet50 v1.5 模型转换为 ONNX 模型并验证正确性，可以在 resnet\_to\_onnx.py 中找到参考代码。 + +#### 如何生成 ONNX 模型 + +**步骤一：将网络权重保存到磁盘** + +首先将训练得到的网络权重保存到磁盘，例如我们保存到 /tmp/resnet50_weights 这个文件夹下 + +```python +check_point = flow.train.CheckPoint() +check_point.save("/tmp/resnet50_weights") +``` + +**步骤二：新建一个用于推理的 job function** + +然后新建一个用于推理的 job function，它只包含网络结构本身，不包含读取 OFRecord 的算子，并且直接接受 numpy 数组形式的输入。可参考 resnet\_to\_onnx.py 中的 `InferenceNet`。 + +**步骤三：调用 flow.onnx.export 方法** + +接下来调用 `flow.onnx.export` 方法，从 OneFlow 网络得到 ONNX 模型，它的第一个参数是上文所说的专用于推理的 job function，第二个参数是 /tmp/resnet50_weights 这个保存了网络权重的文件夹，第三个参数是 ONNX 模型文件的路径。 + +```python +flow.onnx.export(InferenceNet, '/tmp/resnet50_weights', 'resnet50_v1.5.onnx') +``` + +#### 验证 ONNX 模型的正确性 + +生成 ONNX 模型之后可以使用 ONNX Runtime 运行 ONNX 模型，以验证 OneFlow 模型和 ONNX 模型能够在相同的输入下产生相同的结果。相应的代码在 resnet\_to\_onnx.py 的 `check_equality`。 diff --git a/Classification/cnns/resnet_to_onnx.py b/Classification/cnns/resnet_to_onnx.py index 9c3a8b069065d840e9f7c349103d349ad8b05867..a89bbe798ea5c4f45a757b6aa173b4e7d07b2599 100644 --- a/Classification/cnns/resnet_to_onnx.py +++ b/Classification/cnns/resnet_to_onnx.py @@ -2,20 +2,22 @@ # from __future__ import division # from __future__ import print_function +from collections import OrderedDict import os +from PIL import Image +import time +from typing import Callable, Text + +import numpy as np import oneflow as flow +import onnx import onnxruntime as ort -import numpy as np - -import time -from PIL import Image -from collections import OrderedDict from resnet_model import resnet50 from imagenet1000_clsidx_to_labels import clsidx_2_labels -def load_image(image_path): +def load_image(image_path: Text) -> np.ndarray: rgb_mean = [123.68, 116.779, 103.939] rgb_std = [58.393, 57.12, 57.375] print(image_path) @@ -42,74 +44,62 @@ def InferenceNet(images=flow.FixedTensorDef((1, 3, 224, 224))): return predictions -def onnx_inference(image_path, onnx_model_path, ort_optimize=True): +def onnx_inference(image: np.ndarray, onnx_model: onnx.ModelProto): """ test onnx model with onnx runtime - :param image_path: input image path - :param onnx_model_path: path of model.onnx - :param ort_optimize: + :param image: input image, a numpy array + :param onnx_model: onnx model :return: """ - assert os.path.isfile(image_path) and os.path.isfile(onnx_model_path) - ort_sess_opt = ort.SessionOptions() - ort_sess_opt.graph_optimization_level = \ - ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED if ort_optimize else \ - ort.GraphOptimizationLevel.ORT_DISABLE_ALL - sess = ort.InferenceSession(onnx_model_path, sess_options=ort_sess_opt) + assert os.path.isfile(image_path) + sess = ort.InferenceSession(onnx_model.SerializeToString()) assert len(sess.get_outputs()) == 1 and len(sess.get_inputs()) <= 1 ipt_dict = OrderedDict() for ipt in sess.get_inputs(): - ipt_dict[ipt.name] = load_image(image_path) - start = time.time() + ipt_dict[ipt.name] = image onnx_res = sess.run([], ipt_dict)[0] - print('Cost: %.4f s' % (time.time() - start)) - clsidx_onnx = onnx_res.argmax() - print('Onnx >> ', onnx_res.max(), clsidx_2_labels[clsidx_onnx]) + return onnx_res -def oneflow_to_onnx(job_func, flow_model_path, onnx_model_dir, external_data=False): +def oneflow_to_onnx(job_func: Callable, flow_weights_path: Text, onnx_model_dir: Text, external_data: bool=False): """ convert oneflow model to onnx model - :param job_func: inference function in oneflow - :param flow_model_path: input oneflow model path - :param onnx_model_dir: output dir path to save model.onnx - :param external_data: - :return: ture or false + :param job_func: inference function in oneflow + :param flow_weights_path: input oneflow model path + :param onnx_model_dir: output dir path to save model.onnx + :return: onnx model """ if not os.path.exists(onnx_model_dir): os.makedirs(onnx_model_dir) - assert os.path.exists(flow_model_path) and os.path.isdir(onnx_model_dir) + assert os.path.exists(flow_weights_path) and os.path.isdir(onnx_model_dir) - check_point = flow.train.CheckPoint() - # it is a trick to keep check_point.save() from hanging when there is no variable - @flow.global_function(flow.FunctionConfig()) - def add_var(): - return flow.get_variable( - name="trick", - shape=(1,), - dtype=flow.float, - initializer=flow.random_uniform_initializer(), - ) - check_point.init() - - onnx_model_path = os.path.join(onnx_model_dir, os.path.basename(flow_model_path) + '.onnx') - flow.onnx.export(job_func, flow_model_path, onnx_model_path, opset=11, external_data=external_data) + onnx_model_path = os.path.join(onnx_model_dir, os.path.basename(flow_weights_path) + '.onnx') + flow.onnx.export(job_func, flow_weights_path, onnx_model_path, opset=11, external_data=external_data) print('Convert to onnx success! >> ', onnx_model_path) - return onnx_model_path + return onnx.load_model(onnx_model_path) + + +def check_equality(job_func: Callable, onnx_model: onnx.ModelProto, image_path: Text) -> (bool, np.ndarray): + image = load_image(image_path) + onnx_res = onnx_inference(image, onnx_model) + oneflow_res = job_func(image).get().ndarray() + is_equal = np.allclose(onnx_res, oneflow_res, rtol=1e-4, atol=1e-5) + return is_equal, onnx_res if __name__ == "__main__": - # path = 'tiger.jpg' - path = 'test_img/ILSVRC2012_val_00020287.JPEG' - flow_model_path = '/your/oneflow/model/path' + # image_path = 'tiger.jpg' + image_path = 'test_img/ILSVRC2012_val_00020287.JPEG' + flow_weights_path = '/your/oneflow/weights/path' onnx_model_dir = 'onnx/model' - # conver oneflow to onnx - onnx_model_path = oneflow_to_onnx(InferenceNet, flow_model_path, onnx_model_dir, external_data=False) + check_point = flow.train.CheckPoint() + check_point.load(flow_weights_path) - # inference - onnx_inference(path, onnx_model_path) + # conver oneflow to onnx + onnx_model = oneflow_to_onnx(InferenceNet, flow_weights_path, onnx_model_dir, external_data=False) - # Output: - # ILSVRC2012_val_00020287.JPEG - # Cost: 0.0319s - # Onnx >> 0.9924272 hay + # check equality + are_equal, onnx_res = check_equality(InferenceNet, onnx_model, image_path) + clsidx_onnx = onnx_res.argmax() + print('Are the results equal? {}'.format('Yes' if are_equal else 'No')) + print('Class: {}; score: {}'.format(clsidx_2_labels[clsidx_onnx], onnx_res.max()))