提交 a10ad9c2 编写于 作者: C Channingss

update code

上级 cfd80849
...@@ -202,7 +202,7 @@ def paddle2onnx(model_path, save_dir, opset_version=10): ...@@ -202,7 +202,7 @@ def paddle2onnx(model_path, save_dir, opset_version=10):
model.program, model.program,
save_dir, save_dir,
scope=fluid.global_scope(), scope=fluid.global_scope(),
opset_number=opset_number) opset_version=opset_version)
def main(): def main():
...@@ -269,7 +269,7 @@ def main(): ...@@ -269,7 +269,7 @@ def main():
elif args.framework == "paddle2onnx": elif args.framework == "paddle2onnx":
assert args.model is not None, "--model should be defined while translating paddle model to onnx" assert args.model is not None, "--model should be defined while translating paddle model to onnx"
paddle2onnx(args.model, args.save_dir, args.onnx_opset) paddle2onnx(args.model, args.save_dir, opset_version=args.onnx_opset)
else: else:
raise Exception( raise Exception(
......
...@@ -33,8 +33,8 @@ class PaddleOpMapper(object): ...@@ -33,8 +33,8 @@ class PaddleOpMapper(object):
self.name_counter = dict() self.name_counter = dict()
self.op_set = None self.op_set = None
def convert(self, program, save_dir, scope=None, opset_number=10): def convert(self, program, save_dir, scope=None, opset_version=10):
self.op_set = self.create_opset(opset_number) self.op_set = self.create_opset(opset_version)
weight_nodes = self.op_set.convert_weights(program, scope=scope) weight_nodes = self.op_set.convert_weights(program, scope=scope)
op_nodes = list() op_nodes = list()
input_nodes = list() input_nodes = list()
...@@ -77,7 +77,7 @@ class PaddleOpMapper(object): ...@@ -77,7 +77,7 @@ class PaddleOpMapper(object):
initializer=[], initializer=[],
inputs=input_nodes, inputs=input_nodes,
outputs=output_nodes) outputs=output_nodes)
opset_imports = [helper.make_opsetid("", opset_number)] opset_imports = [helper.make_opsetid("", opset_version)]
model = helper.make_model( model = helper.make_model(
graph, producer_name='X2Paddle', opset_imports=opset_imports) graph, producer_name='X2Paddle', opset_imports=opset_imports)
onnx.checker.check_model(model) onnx.checker.check_model(model)
...@@ -89,20 +89,20 @@ class PaddleOpMapper(object): ...@@ -89,20 +89,20 @@ class PaddleOpMapper(object):
print("\nTranslated model saved in {}".format( print("\nTranslated model saved in {}".format(
os.path.join(save_dir, 'x2paddle_model.onnx'))) os.path.join(save_dir, 'x2paddle_model.onnx')))
def create_opset(self, opset_number): def create_opset(self, opset_version=10):
run_opset = self.default_opset run_opset = self.default_opset
opset = '' opset = ''
if opset_number in self.support_opsets: if opset_version in self.support_opsets:
run_opset = opset_number run_opset = opset_version
else: else:
for support_opset_number in self.support_opsets: for support_opset_version in self.support_opsets:
if support_opset_number < opset_number: if support_opset_version < opset_version:
run_opset = support_opset_number run_opset = support_opset_version
else: else:
break break
print( print(
'Now, onnx2paddle support convert onnx model opset_verison {},' 'Now, onnx2paddle support convert onnx model opset_verison {},'
'opset_verison of your onnx model is {}, automatically treated as op_set: {}.' 'opset_verison of your onnx model is {}, automatically treated as op_set: {}.'
.format(self.support_opsets, opset_number, run_opset)) .format(self.support_opsets, opset_version, run_opset))
opset = 'OpSet' + str(run_opset) opset = 'OpSet' + str(run_opset)
return eval(opset)() return eval(opset)()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册