1
1
"""
2
2
Copyright 2020 The OneFlow Authors. All rights reserved.
3
-
4
3
Licensed under the Apache License, Version 2.0 (the "License");
5
4
you may not use this file except in compliance with the License.
6
5
You may obtain a copy of the License at
7
-
8
6
http://www.apache.org/licenses/LICENSE-2.0
9
-
10
7
Unless required by applicable law or agreed to in writing, software
11
8
distributed under the License is distributed on an "AS IS" BASIS,
12
9
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
29
26
from resnet_model import resnet50
30
27
import config as configs
31
28
from imagenet1000_clsidx_to_labels import clsidx_2_labels
29
+ from oneflow_onnx .oneflow2onnx .util import export_onnx_model
32
30
33
31
parser = configs .get_parser ()
34
32
args = parser .parse_args ()
@@ -92,12 +90,12 @@ def oneflow_to_onnx(
92
90
assert os .path .exists (flow_weights_path ) and os .path .isdir (onnx_model_dir )
93
91
94
92
onnx_model_path = os .path .join (
95
- onnx_model_dir , os . path . basename ( flow_weights_path ) + " .onnx"
93
+ onnx_model_dir , "model .onnx"
96
94
)
97
- flow . onnx . export (
95
+ export_onnx_model (
98
96
job_func ,
99
- flow_weights_path ,
100
- onnx_model_path ,
97
+ flow_weight_dir = flow_weights_path ,
98
+ onnx_model_path = onnx_model_dir ,
101
99
opset = 11 ,
102
100
external_data = external_data ,
103
101
)
@@ -132,4 +130,4 @@ def check_equality(
132
130
are_equal , onnx_res = check_equality (InferenceNet , onnx_model , image_path )
133
131
clsidx_onnx = onnx_res .argmax ()
134
132
print ("Are the results equal? {}" .format ("Yes" if are_equal else "No" ))
135
- print ("Class: {}; score: {}" .format (clsidx_2_labels [clsidx_onnx ], onnx_res .max ()))
133
+ print ("Class: {}; score: {}" .format (clsidx_2_labels [clsidx_onnx ], onnx_res .max ()))
0 commit comments