11"""
22Copyright 2020 The OneFlow Authors. All rights reserved.
3-
43Licensed under the Apache License, Version 2.0 (the "License");
54you may not use this file except in compliance with the License.
65You may obtain a copy of the License at
7-
86 http://www.apache.org/licenses/LICENSE-2.0
9-
107Unless required by applicable law or agreed to in writing, software
118distributed under the License is distributed on an "AS IS" BASIS,
129WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
2926from resnet_model import resnet50
3027import config as configs
3128from imagenet1000_clsidx_to_labels import clsidx_2_labels
29+ from oneflow_onnx .oneflow2onnx .util import export_onnx_model
3230
3331parser = configs .get_parser ()
3432args = parser .parse_args ()
@@ -92,12 +90,12 @@ def oneflow_to_onnx(
9290 assert os .path .exists (flow_weights_path ) and os .path .isdir (onnx_model_dir )
9391
9492 onnx_model_path = os .path .join (
95- onnx_model_dir , os . path . basename ( flow_weights_path ) + " .onnx"
93+ onnx_model_dir , "model .onnx"
9694 )
97- flow . onnx . export (
95+ export_onnx_model (
9896 job_func ,
99- flow_weights_path ,
100- onnx_model_path ,
97+ flow_weight_dir = flow_weights_path ,
98+ onnx_model_path = onnx_model_dir ,
10199 opset = 11 ,
102100 external_data = external_data ,
103101 )
@@ -132,4 +130,4 @@ def check_equality(
132130 are_equal , onnx_res = check_equality (InferenceNet , onnx_model , image_path )
133131 clsidx_onnx = onnx_res .argmax ()
134132 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