diff --git a/Classification/cnns/README.md b/Classification/cnns/README.md index 4e63671..56c4465 100644 --- a/Classification/cnns/README.md +++ b/Classification/cnns/README.md @@ -586,9 +586,9 @@ onnx_model_dir = 'onnx/model' **步骤三:调用 flow.onnx.export 方法** -接下来代码中会调用`oneflow_to_onnx()`方法,此方法包含了核心的模型转换方法: `flow.onnx.export()` +接下来代码中会调用`oneflow_to_onnx()`方法,此方法包含了核心的模型转换方法: `oneflow_onnx.oneflow2onnx.util.export_onnx_model()`,更多OneFlow和ONNX模型转换相关的问题请看: [oneflow_convert_tools介绍](https://docs.oneflow.org/extended_topics/oneflow_convert_tools.html) - **flow.onnx.export** 将从 OneFlow 网络得到 ONNX 模型,它的第一个参数是上文所说的专用于推理的 job function,第二个参数是OneFlow模型路径,第三个参数是(转换后)ONNX模型的存放路径 + **oneflow_to_onnx** 将从 OneFlow 网络得到 ONNX 模型,它的第一个参数是上文所说的专用于推理的 job function,第二个参数是OneFlow模型路径,第三个参数是(转换后)ONNX模型的存放路径 ```python onnx_model = oneflow_to_onnx(InferenceNet, flow_weights_path, onnx_model_dir, external_data=False) diff --git a/Classification/cnns/resnet_to_onnx.py b/Classification/cnns/resnet_to_onnx.py index 1af0344..a0a0fb8 100644 --- a/Classification/cnns/resnet_to_onnx.py +++ b/Classification/cnns/resnet_to_onnx.py @@ -1,12 +1,9 @@ """ Copyright 2020 The OneFlow Authors. All rights reserved. - 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. @@ -29,6 +26,7 @@ from resnet_model import resnet50 import config as configs from imagenet1000_clsidx_to_labels import clsidx_2_labels +from oneflow_onnx.oneflow2onnx.util import export_onnx_model parser = configs.get_parser() args = parser.parse_args() @@ -92,12 +90,12 @@ def oneflow_to_onnx( assert os.path.exists(flow_weights_path) and os.path.isdir(onnx_model_dir) onnx_model_path = os.path.join( - onnx_model_dir, os.path.basename(flow_weights_path) + ".onnx" + onnx_model_dir, "model.onnx" ) - flow.onnx.export( + export_onnx_model( job_func, - flow_weights_path, - onnx_model_path, + flow_weight_dir=flow_weights_path, + onnx_model_path=onnx_model_dir, opset=11, external_data=external_data, ) @@ -132,4 +130,4 @@ def 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())) + print("Class: {}; score: {}".format(clsidx_2_labels[clsidx_onnx], onnx_res.max())) \ No newline at end of file