def _test(self, model, inputShapes, testedOutputs=[]):
modelName = re.compile("^.*?/?([^/.]+).prototext$").search(model).group(1)
o, miniBatchSize = lbann.onnx.l2o.parseLbannModelPB(model, inputShapes)
if ADD_DUMMY_PARAMS:
dummyInits = []
for i in o.graph.input:
if i.name in list(map(lambda x: "{}_0".format(x), inputShapes.keys())) or \
i.name in list(map(lambda x: x.name, o.graph.initializer)):
continue
shape = lbann.onnx.util.getDimFromValueInfo(i)
dummyInits.append(numpy_helper.from_array(np.zeros(shape, dtype=lbann.onnx.ELEM_TYPE_NP),
name=i.name))
g = onnx.helper.make_graph(o.graph.node,
o.graph.name,
o.graph.input,
o.graph.output,
list(o.graph.initializer) + dummyInits,
value_info=o.graph.value_info)
o = onnx.helper.make_model(g)
if SAVE_ONNX:
onnx.save(o, os.path.join(DUMP_DIR, "{}.onnx".format(modelName)))
for nodeName, outputShape in testedOutputs:
node, = list(filter(lambda x: x.name == nodeName, o.graph.node))
outputVI, = list(filter(lambda x: x.name == node.output[0], o.graph.value_info))
outputShapeActual = lbann.onnx.util.getDimFromValueInfo(outputVI)
outputShapeWithMB = list(map(lambda x: miniBatchSize if x == MB_PLACEHOLDER else x, outputShape))
assert outputShapeWithMB == outputShapeActual, (outputShapeWithMB, outputShapeActual)
def main():
xmodel = onnx.load(sys.argv[1])
subgraphs = find_subgraphs(xmodel.graph)
if len(sys.argv) == 2:
for name, graph in subgraphs.items():
print('name=%s node=%d' % (name, len(graph.node)))
else:
g = subgraphs[sys.argv[2]]
m = onnx.helper.make_model(g)
onnx.save(m, 'out.onnx')
def main():
from_dir = sys.argv[1]
to_dir = sys.argv[2]
os.makedirs(to_dir, exist_ok=True)
xmodel = onnx.load(os.path.join(from_dir, 'model.onnx'))
convert_model(xmodel)
onnx.save(xmodel, os.path.join(to_dir, 'model.onnx'))
for test_dir in glob.glob(os.path.join(from_dir, 'test_data_set_*')):
dir_name = os.path.basename(test_dir)
to_test_dir = os.path.join(to_dir, dir_name)
os.makedirs(to_test_dir, exist_ok=True)
for pb_filename in glob.glob(os.path.join(test_dir, '*.pb')):
pb_name = os.path.basename(pb_filename)
tensor = onnx.load_tensor(pb_filename)
convert_tensor(tensor)
onnx.save_tensor(tensor, os.path.join(to_test_dir, pb_name))
def export(
self,
output: str,
input_example=None,
output_example=None,
verbose=False,
export_params=True,
do_constant_folding=True,
keep_initializers_as_inputs=False,
onnx_opset_version: int = 12,
try_script: bool = False,
set_eval: bool = True,
check_trace: bool = True,
use_dynamic_axes: bool = True,
):
if input_example is not None or output_example is not None:
logging.warning(
"Passed input and output examples will be ignored and recomputed since"
" TokenClassificationModel consists of two separate models with different"
" inputs and outputs.")
qual_name = self.__module__ + '.' + self.__class__.__qualname__
output1 = os.path.join(os.path.dirname(output),
'bert_' + os.path.basename(output))
output1_descr = qual_name + ' BERT exported to ONNX'
bert_model_onnx = self.bert_model.export(
output1,
None, # computed by input_example()
None,
verbose,
export_params,
do_constant_folding,
keep_initializers_as_inputs,
onnx_opset_version,
try_script,
set_eval,
check_trace,
use_dynamic_axes,
)
output2 = os.path.join(os.path.dirname(output),
'classifier_' + os.path.basename(output))
output2_descr = qual_name + ' Classifier exported to ONNX'
classifier_onnx = self.classifier.export(
output2,
None, # computed by input_example()
None,
verbose,
export_params,
do_constant_folding,
keep_initializers_as_inputs,
onnx_opset_version,
try_script,
set_eval,
check_trace,
use_dynamic_axes,
)
output_model = attach_onnx_to_onnx(bert_model_onnx, classifier_onnx,
"TKCL")
output_descr = qual_name + ' BERT+Classifier exported to ONNX'
onnx.save(output_model, output)
return ([output, output1,
output2], [output_descr, output1_descr, output2_descr])
# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
import onnx
from onnx import helper, TensorProto
INPUT_1 = helper.make_tensor_value_info('foo', TensorProto.FLOAT, [1])
INPUT_2 = helper.make_tensor_value_info('baz', TensorProto.FLOAT, [1])
OUTPUT = helper.make_tensor_value_info('bar', TensorProto.FLOAT, [1])
nodes = [
helper.make_node(
'Mul',
['foo', 'baz'],
['bat'],
),
]
graph_def = helper.make_graph(
nodes,
'mul',
[INPUT_1, INPUT_2],
[OUTPUT],
)
model_def = helper.make_model(
graph_def,
producer_name='foobar.py',
opset_imports=[onnx.OperatorSetIdProto(version=12)])
onnx.save(model_def, 'foobar.onnx')
def main():
# Configurable parameters from command line
parser = argparse.ArgumentParser(description='ONNX Modifying Example')
parser.add_argument('--onnx', help='onnx file to modify')
parser.add_argument(
'--output',
default="output.onnx",
help='input batch size for testing (default: output.onnx)')
args = parser.parse_args()
# Load ONNX file
model = onnx.load(args.onnx)
# Retrieve graph_def
graph = model.graph
node_input_new = False
counter_conv_nodes_updated = 0
nodes_to_delete = []
# Iterate through all the nodes
for i, node in enumerate(graph.node):
if not node_input_new:
node_input_new = graph.node[0].input[0]
if counter_conv_nodes_updated == 2:
break
if node.op_type == 'Conv':
# Update inputs of any Conv node and converting Conv->CoordConv
graph.node[i].input.remove(graph.node[i].input[0])
graph.node[i].input.insert(0, node_input_new)
graph.node[i].op_type = COORD_CONV_OP_TYPE
counter_conv_nodes_updated += 1
elif node.op_type == 'Relu':
# Saving output of previous node
node_input_new = graph.node[i].output[0]
else:
# Add node to list of removable nodes
nodes_to_delete.append(i)
for i in nodes_to_delete[::-1]:
# Remove unnecessary nodes
n = graph.node[i]
graph.node.remove(n)
# insert AC nodes
i = 0
while i < len(graph.node):
if graph.node[i].op_type == COORD_CONV_OP_TYPE:
print('here')
# Create an ac node
node_ac = onnx.NodeProto()
node_ac.op_type = "CoordConvAC"
node_ac.output.insert(0, f"ac_output_{i}")
node_ac.input.insert(0, graph.node[i].input[0])
graph.node[i].input[0] = f"ac_output_{i}"
graph.node.insert(i, node_ac)
i += 1
i += 1
# Generate model_cropped from modified graph
model_cropped = onnx.helper.make_model(graph)
print(onnx.helper.printable_graph(model_cropped.graph))
print("Inputs:", model_cropped.graph.node[0].input, "Outputs:",
model_cropped.graph.node[-1].output)
# Save the serialized model
onnx.save(model_cropped, args.output)
def generate_gemm_scan_model(model_name, config1, config2):
model = onnx.ModelProto()
model.ir_version = 7 # use stable onnx ir version
opset = model.opset_import.add()
opset.version = 11
# Based on the given configs, we would have a model like below:
# Main graph, where C is an initializer and passed as the input state for the Scan:
# C input_1A input_2A
# \ | /
# \ | /
# Scan
# |
# output
#
# Scan's subgraph, where out_C is the output state of the Scan
# input_1A B C input_2A B C
# \ | / \ | /
# \ |/ \ |/
# Gemm_1 Gemm_2
# \ /
# \ /
# Sub
# / \
# out_C output
#
# config1 and config2 configure alpha/beta/transA/transB for Gemm_1 and Gemm_2, respectively.
scan_body = onnx.GraphProto()
scan_body.name = 'gemm_subgraph'
shape_c1 = [config1['M'], config1['N']]
shape_c2 = [config2['M'], config2['N']]
assert shape_c1 == shape_c2
C1 = config1['C']
C2 = config2['C']
scan_node_inputs = []
postfix = '_subgraph'
states_cnt = 0
# make sure we create state inputs first
if config1['withC']:
assert config1['initC']
states_cnt = states_cnt + 1
scan_node_inputs.append(C1)
scan_body.input.add().CopyFrom(
helper.make_tensor_value_info('in_' + C1 + postfix,
onnx.TensorProto.FLOAT, shape_c1))
if config2['withC'] and C1 != C2:
assert config2['initC']
states_cnt = states_cnt + 1
scan_node_inputs.append(C2)
scan_body.input.add().CopyFrom(
helper.make_tensor_value_info('in_' + C2 + postfix,
onnx.TensorProto.FLOAT, shape_c2))
added_inputs_subgraph = {}
generate_gemm_node_subgraph(scan_body, scan_node_inputs, postfix, config1,
added_inputs_subgraph)
generate_gemm_node_subgraph(scan_body, scan_node_inputs, postfix, config2,
added_inputs_subgraph)
sub_output = 'sub_output' + postfix
# create a Sub op instead of Add to break the MatMul-to-Gemm rewriting rule
# performed by the ort optimizer
sub_node = helper.make_node(
'Sub', [config1['Y'] + postfix, config2['Y'] + postfix], [sub_output],
'sub_node')
scan_body.node.add().CopyFrom(sub_node)
scan_node_outputs = []
# create state outputs
if config1['withC']:
id_node1 = onnx.helper.make_node('Identity', [sub_output],
['out_' + C1 + postfix], 'id_node1')
scan_body.node.add().CopyFrom(id_node1)
scan_body.output.add().CopyFrom(
helper.make_tensor_value_info('out_' + C1 + postfix,
onnx.TensorProto.FLOAT, shape_c1))
scan_node_outputs.append('out_' + C1)
if config2['withC'] and C1 != C2:
id_node2 = onnx.helper.make_node('Identity', [sub_output],
['out_' + C2 + postfix], 'id_node2')
scan_body.node.add().CopyFrom(id_node2)
scan_body.output.add().CopyFrom(
helper.make_tensor_value_info('out_' + C2 + postfix,
onnx.TensorProto.FLOAT, shape_c2))
scan_node_outputs.append('out_' + C2)
# scan subgraph output
scan_body.output.add().CopyFrom(
helper.make_tensor_value_info(sub_output, onnx.TensorProto.FLOAT,
shape_c1))
scan_node_outputs.append('scan_output')
# create scan node
inputs_cnt = len(scan_node_inputs) - states_cnt
assert inputs_cnt > 0
scan_node = onnx.helper.make_node('Scan',
scan_node_inputs,
scan_node_outputs,
'scan_node',
num_scan_inputs=inputs_cnt,
body=scan_body)
model.graph.node.add().CopyFrom(scan_node)
added_inputs_initializers = {}
# main graph inputs and initializers
(a1, b1, c1) = generate_gemm_inputs_initializers(model.graph,
config1,
added_inputs_initializers,
extend=True)
(a2, b2, c2) = generate_gemm_inputs_initializers(model.graph,
config2,
added_inputs_initializers,
extend=True)
shape_output = ['seq', config1['M'], config1['N']]
# main graph outputs
model.graph.output.add().CopyFrom(
helper.make_tensor_value_info('scan_output', onnx.TensorProto.FLOAT,
shape_output))
onnx.save(model, model_name)
return (a1, b1, c1, a2, b2, c2)
def generate_qat_model(model_names):
test_models = []
test_initializers = []
'''
TEST_MODEL_CONFIG_1
'''
# Main graph:
#
# [A] [input_bias]
# \ /
# Add [scale_zp_const] [input_weight]
# | \ /
# | QuantizeLinear_1
# QuantizeLinear_0 |
# | DequantizeLinear_1
# | /
# DequantizeLinear_0 Transpose
# \ /
# \ / <--- (actual graph: this branch is folded)
# Matmul
# |
# |
# [B]
graph = helper.make_graph(
[
#nodes
helper.make_node("Add", ["A", "input_bias"], ["add_out"], "add0"),
helper.make_node(
"QuantizeLinear",
["add_out", "quant0_scale_const", "quant0_zp_const"],
["quant0_out"], "qlinear0"),
helper.make_node(
"DequantizeLinear",
["quant0_out", "dequant0_scale_const", "dequant0_zp_const"],
["dequant0_out"], "dqlinear0"),
helper.make_node("MatMul", ["dequant0_out", "trans_out"], ["B"],
"matmul"),
],
"QAT_model_1", #name
[ #input
helper.make_tensor_value_info('A', TensorProto.FLOAT, ['unk_1'])
],
[ #output
helper.make_tensor_value_info('B', TensorProto.FLOAT, [1024])
],
[ #initializers
helper.make_tensor('quant0_scale_const', TensorProto.FLOAT, [],
[0.01961481384932995]),
helper.make_tensor('quant0_zp_const', TensorProto.INT8, [], [0]),
helper.make_tensor('dequant0_scale_const', TensorProto.FLOAT, [],
[0.01961481384932995]),
helper.make_tensor('dequant0_zp_const', TensorProto.INT8, [], [0]),
])
input_weight_1 = generate_input_initializer([1024, 1024], np.float32,
'trans_out')
input_bias_1 = generate_input_initializer([1024], np.float32, 'input_bias')
graph.initializer.add().CopyFrom(input_weight_1)
graph.initializer.add().CopyFrom(input_bias_1)
model_1 = onnx.helper.make_model(graph)
model_1.ir_version = onnx.IR_VERSION
onnx.save(model_1, model_names[0])
test_models.extend([model_1])
initiazliers_1 = [input_weight_1, input_bias_1]
test_initializers.append(initiazliers_1)
'''
TEST_MODEL_CONFIG_2
'''
# Main graph:
#
# [A]
# |
# MaxPool
# / \
# QuantizeLinear_0 QuantizeLinear_1
# | |
# DequantizeLinear_0 DequantizeLinear_1
# | |
# Conv_0-[weight,bias] Conv_1-[weight,bias]
# \ /
# \ /
# Add
# |
# [B]
graph = helper.make_graph(
[
#nodes
helper.make_node("MaxPool", ["A"], ["maxpool_out"], "maxpool"),
helper.make_node(
"QuantizeLinear",
["maxpool_out", "quant0_scale_const", "quant0_zp_const"],
["quant0_out"], "qlinear0"),
helper.make_node(
"DequantizeLinear",
["quant0_out", "dequant0_scale_const", "dequant0_zp_const"],
["dequant0_out"], "dqlinear0"),
helper.make_node("Conv", ["dequant0_out"], ["conv0_out"], "conv0"),
helper.make_node(
"QuantizeLinear",
["maxpool_out", "quant1_scale_const", "quant1_zp_const"],
["quant1_out"], "qlinear1"),
helper.make_node(
"DequantizeLinear",
["quant1_out", "dequant1_scale_const", "dequant1_zp_const"],
["dequant1_out"], "dqlinear1"),
helper.make_node("Conv", ["dequant1_out"], ["conv1_out"], "conv1"),
helper.make_node("Add", ["conv0_out", "conv1_out"], ["B"], "add"),
],
"QAT_model_2", #name
[ #input
helper.make_tensor_value_info('A', TensorProto.FLOAT, ['unk_1'])
],
[ #output
helper.make_tensor_value_info('B', TensorProto.FLOAT,
[256, 64, 1, 1])
],
[ #initializers
helper.make_tensor('quant0_scale_const', TensorProto.FLOAT, [],
[0.2062656134366989]),
helper.make_tensor('quant0_zp_const', TensorProto.UINT8, [],
[165]),
helper.make_tensor('dequant0_scale_const', TensorProto.FLOAT, [],
[0.2062656134366989]),
helper.make_tensor('dequant0_zp_const', TensorProto.UINT8, [],
[165]),
helper.make_tensor('quant1_scale_const', TensorProto.FLOAT, [],
[0.10088317096233368]),
helper.make_tensor('quant1_zp_const', TensorProto.UINT8, [],
[132]),
helper.make_tensor('dequant1_scale_const', TensorProto.FLOAT, [],
[0.10088317096233368]),
helper.make_tensor('dequant1_zp_const', TensorProto.UINT8, [],
[132]),
])
conv_weight_0 = generate_input_initializer([256, 64, 1, 1], np.float32,
'conv_weight_0')
conv_bias_0 = generate_input_initializer([256], np.float32, 'conv_bias_0')
graph.initializer.add().CopyFrom(conv_weight_0)
graph.initializer.add().CopyFrom(conv_bias_0)
conv_weight_1 = generate_input_initializer([256, 64, 1, 1], np.float32,
'conv_weight_1')
conv_bias_1 = generate_input_initializer([256], np.float32, 'conv_bias_1')
graph.initializer.add().CopyFrom(conv_weight_1)
graph.initializer.add().CopyFrom(conv_bias_1)
model_2 = onnx.helper.make_model(graph)
model_2.ir_version = onnx.IR_VERSION
onnx.save(model_2, model_names[1])
test_models.extend([model_2])
initializers_2 = [conv_weight_0, conv_bias_0, conv_weight_1, conv_weight_1]
test_initializers.append(initializers_2)
return test_models, test_initializers
def export_onnx(model,
im,
file,
opset,
train,
dynamic,
simplify,
prefix=colorstr('ONNX:')):
# YOLOv5 ONNX export
try:
check_requirements(('onnx', ))
import onnx
print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
f = file.with_suffix('.onnx')
torch.onnx.export(
model,
im,
f,
verbose=False,
opset_version=opset,
training=torch.onnx.TrainingMode.TRAINING
if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train,
input_names=['images'],
output_names=['output'],
dynamic_axes={
'images': {
0: 'batch',
2: 'height',
3: 'width'
}, # shape(1,3,640,640)
'output': {
0: 'batch',
1: 'anchors'
} # shape(1,25200,85)
} if dynamic else None)
# Checks
model_onnx = onnx.load(f) # load onnx model
onnx.checker.check_model(model_onnx) # check onnx model
# print(onnx.helper.printable_graph(model_onnx.graph)) # print
# Simplify
if simplify:
try:
check_requirements(('onnx-simplifier', ))
import onnxsim
print(
f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...'
)
model_onnx, check = onnxsim.simplify(
model_onnx,
dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)}
if dynamic else None)
assert check, 'assert check failed'
onnx.save(model_onnx, f)
except Exception as e:
print(f'{prefix} simplifier failure: {e}')
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
print(
f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'"
)
except Exception as e:
print(f'{prefix} export failure: {e}')
[1, 1])
OUTPUT = helper.make_tensor_value_info('output', TensorProto.INT32, [1, 1])
nodes = [
helper.make_node(
'MatMul',
['query_tensor', 'attribute_tensor'],
['matmul'],
),
helper.make_node(
'Add',
['matmul', 'bias_tensor'],
['output'],
),
]
graph_def = helper.make_graph(
nodes,
'int_types_scoring',
[
QUERY_TENSOR,
ATTRIBUTE_TENSOR,
BIAS_TENSOR,
],
[OUTPUT],
)
model_def = helper.make_model(
graph_def,
producer_name='int_types.py',
opset_imports=[onnx.OperatorSetIdProto(version=12)])
onnx.save(model_def, 'int_types.onnx')
def convert(onnx_model_filename,
save_dir,
model_basename='model.py',
model_func_name='inference',
embed_params=False,
onnx_opset_version=9,
onnx_opset_pedantic=True,
debug=False):
"""
convert an ONNX model to Paddle fluid Python code and desc pb
"""
import onnx
from onnx.checker import ValidationError
from onnx.checker import check_model
from onnx.utils import polish_model
from onnx.version_converter import convert_version
try:
from . import onnx_utils, writer
except ImportError:
import onnx_utils, writer
# imports
DEFAULT_OP_DOMAIN = onnx_utils.DEFAULT_OP_DOMAIN
graph_ops, graph_weights = onnx_utils.graph_ops, onnx_utils.graph_weights
inferred_model_value_info = onnx_utils.inferred_model_value_info
optimize_model_skip_op_for_inference = onnx_utils.optimize_model_skip_op_for_inference
optimize_model_strip_initializer = onnx_utils.optimize_model_strip_initializer
optimize_model_cast = onnx_utils.optimize_model_cast
optimize_model_slice = onnx_utils.optimize_model_slice
Program, Writer = writer.Program, writer.Writer
make_var_name = writer.make_var_name
logger = logging.getLogger('convert')
# prepare onnx model
logger.info('loading model: %s ...', onnx_model_filename)
onnx_model = onnx.load(onnx_model_filename)
try:
logger.info('checking model ...')
check_model(onnx_model)
logger.debug('using opset version: %d', onnx_opset_version)
if onnx_opset_pedantic: # WORKAROUND: RuntimeError: No Adapter For OP
onnx_model = convert_version(onnx_model, onnx_opset_version)
else: # TODO: add new argument for this option
logger.warning(
'opset conversion skipped for onnx_opset_pedantic is OFF')
onnx_model = polish_model(onnx_model)
except ValidationError as e:
if onnx_opset_pedantic:
raise e
else:
logger.warning('due to onnx_opset_pedantic is OFF')
logger.warning('the ONNX model sanity checking error is suppressed')
logger.warning('value_info inferring may be uncompleted')
# onnx model optimization
logger.info('optimizing model ...')
onnx_model = optimize_model_skip_op_for_inference(onnx_model)
onnx_model = optimize_model_strip_initializer(onnx_model)
onnx_model = optimize_model_cast(onnx_model)
onnx_model = optimize_model_slice(onnx_model)
# prepare filesystem
shutil.rmtree(save_dir, ignore_errors=True)
shutil.os.makedirs(save_dir, exist_ok=True)
logger.info('folder %s cleared', save_dir)
# DEBUG:
if debug:
model = onnx.shape_inference.infer_shapes(onnx_model)
debug_model_filename, _ = shutil.os.path.splitext(onnx_model_filename)
onnx.save(model, debug_model_filename + '.optimized_and_inffered.onnx')
# onnx.save(model, '/tmp/export/optimized_and_inffered.onnx')
# I/O instances
onnx_graph = onnx_model.graph
fluid_program = Program()
fluid_writer = Writer()
# model components
# graph_name = onnx_graph.name
graph_inputs = [value.name for value in onnx_graph.input]
graph_outputs = [value.name for value in onnx_graph.output]
graph_params = []
graph_value_infos = inferred_model_value_info(onnx_model)
# prepare additional value_info
# for weights
for name, weight in graph_weights(onnx_graph):
value_info = graph_value_infos[name]
value_info['embeded_as'] = []
value_info['get_weight'] = (lambda w: lambda: w.tolist())(
weight) # lazy getter
logger.info('conversion started')
# op set conversion
# topo = 'backward' if embed_params else 'forward'
topo = 'forward'
for name, domain, op_type, inputs, outputs, attrs in graph_ops(
onnx_graph, topo=topo):
logger.debug('translating op %s %s::%s ...', name, domain, op_type)
if domain == DEFAULT_OP_DOMAIN:
domain = ''
try:
fluid_writer.emit_op(
fluid_program,
name,
domain,
op_type,
inputs,
outputs,
attrs,
graph_value_infos,
embed_params=embed_params,
)
except BaseException as e:
logger.fatal('conversion failed for:\n\t%s -> %s::%s -> %s', inputs,
domain, op_type, outputs)
raise e
op_codes = fluid_program.codes
fluid_program.codes = []
logger.info('%d ops converted', len(fluid_program.op_descs))
# weight writer
for name, weight in graph_weights(onnx_graph):
graph_params.append(name)
value_info = graph_value_infos[name]
var_names = value_info.get('embeded_as', [])
if var_names:
if len(var_names) > 1:
logger.info(
'weight %s is shared between ops, more disk space will be consumed',
name)
logger.debug(
'saving weight %s with size of %d, in %d bytes, as %s ...',
name, weight.size, weight.nbytes, var_names)
for var_name in var_names: # multiple references
fluid_writer.write_weight(
weight, shutil.os.path.join(save_dir, var_name))
else:
logger.debug(
'saving weight %s with size of %d, in %d bytes, to %s ...',
name, weight.size, weight.nbytes, make_var_name(name))
fluid_writer.write_weight(
weight, shutil.os.path.join(save_dir, make_var_name(name)))
fluid_writer.emit_param(fluid_program, name, value_info)
param_codes = fluid_program.codes
fluid_program.codes = []
logger.info('%d weights converted', len(graph_params))
# input writer
external_inputs = []
for name in graph_inputs:
if name not in graph_params:
value_info = graph_value_infos[name]
assert value_info['external']
external_inputs.append(name)
fluid_writer.emit_inputs(
fluid_program, external_inputs, graph_value_infos,
remove_batch=False) # TODO:
input_codes = fluid_program.codes
fluid_program.codes = []
logger.info('%d inputs converted', len(external_inputs))
# output writer
external_outputs = []
for name in graph_outputs:
if name not in graph_params:
value_info = graph_value_infos[name]
assert value_info['external']
external_outputs.append(name)
fluid_writer.emit_outputs(fluid_program, external_outputs)
output_codes = [''] + fluid_program.codes # add an empty line
fluid_program.codes = []
logger.info('%d outputs converted', len(external_outputs))
# code generation
header_codes = fluid_writer.header_code(
model_func_name, 'From: {}'.format(onnx_model_filename))
code_filename = shutil.os.path.join(save_dir, model_basename)
fluid_writer.write_code_file(code_filename, header_codes, input_codes,
param_codes, op_codes, output_codes)
logger.info('code saved to %s, factory function: %s', code_filename,
model_func_name)
# desc generation
desc_filename = shutil.os.path.join(save_dir, '__model__')
fluid_writer.write_desc_file(
desc_filename,
op_descs=fluid_program.op_descs,
var_descs=fluid_program.var_descs,
)
logger.info('program saved to %s', desc_filename)
logger.info('conversion finished')
# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
import onnx
import numpy as np
from onnx import helper, TensorProto
INPUT_1 = helper.make_tensor_value_info('input1', TensorProto.FLOAT,
["batch", 2])
INPUT_2 = helper.make_tensor_value_info('input2', TensorProto.FLOAT,
["batch", 2])
OUTPUT = helper.make_tensor_value_info('out', TensorProto.FLOAT,
["batch", "dim1", "dim2"])
SHAPE = helper.make_tensor('shape',
TensorProto.INT64,
dims=[3],
vals=np.array([1, 2, 2]).astype(np.int64))
nodes = [
helper.make_node('Concat', ['input1', 'input2'], ['concat'], axis=1),
helper.make_node('Reshape', ['concat', 'shape'], ['out']),
]
graph_def = helper.make_graph(nodes, 'simple_scoring', [INPUT_1, INPUT_2],
[OUTPUT], [SHAPE])
model_def = helper.make_model(
graph_def,
producer_name='create_dynamic_model.py',
opset_imports=[onnx.OperatorSetIdProto(version=12)])
onnx.save(model_def, 'dynamic_model_2.onnx')
from onnxsim import simplify import onnx import argparse import sys if __name__=='__main__': onnx_model = onnx.load(sys.argv[1]) # load onnx model model_simp, check = simplify(onnx_model) assert check, "Simplified ONNX model could not be validated" onnx.save(model_simp, sys.argv[2])
import sys
import onnx
import onnx.shape_inference
if len(sys.argv) != 3:
raise RuntimeError('Usage: %s input.onnx output.onnx' % sys.argv[0])
model = onnx.load(sys.argv[1], 'rb')
model = onnx.shape_inference.infer_shapes(model)
onnx.save(model, sys.argv[2])
input_names=["input_img"],
output_names=["output"],
dynamic_axes={"input_img": [2, 3]})
sess = onnxruntime.InferenceSession(
'../../onnxTestModels/loopWithConv.onnx')
print('Load model')
result = sess.run(None, {'input_img': img_tensor_numpy})
model = onnx.load('../../onnxTestModels/loopWithConv.onnx')
# check model
onnx.checker.check_model(model)
print("onnx checker done")
# override input shapes
model.graph.input[0].type.tensor_type.shape.dim[2].dim_value = 48
model.graph.input[0].type.tensor_type.shape.dim[3].dim_value = 48
onnx.save(model, '../../onnxTestModels/loopWithConv_override.onnx')
# using ONNX shape inference
inferModel = onnx.shape_inference.infer_shapes(model)
onnx.save(inferModel,
'../../onnxTestModels/loopWithConv_override_inferred.onnx')
# using symbolic shape inference
out_mp = sym.SymbolicShapeInference.infer_shapes(in_mp=model,
auto_merge=True)
onnx.save(out_mp,
'../../onnxTestModels/loopWithConv_override_symbolic.onnx')
model = onnx.load("net.onnx")
graph = model.graph
#Generate a name for all node if they have none.
nodeIdx = 0
for n in graph.node:
print("node: \n\t" + str(n).replace("\n", "\n\t"))
if n.name == '':
n.name = str(n.op_type) + str(nodeIdx)
nodeIdx += 1
#print("name: " + str(n.name))
node_map = createGraphMemberMap(graph.node) # convert to dict
input_map = createGraphMemberMap(graph.input) # convert to dict
output_map = createGraphMemberMap(graph.output) # convert to dict
initializer_map = createGraphMemberMap(graph.initializer) # convert to dict
#print("input_map: " + str(input_map.keys()))
#print("output_map: " + str(output_map.keys()))
#print("node_map: " + str(node_map.keys()))
#print("initializer map: " + str(initializer_map.keys()))
for key in node_map.keys():
print("key: " + key)
if "Clip" in key:
print("removing key: " + str(key))
graph.node.remove(node_map[key])
onnx.save(model, "netout.onnx")
import onnxruntime
from onnxsim import simplify
import argparse
import os
from pprint import pprint
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, required=True)
parser.add_argument("--height", type=int, required=True)
parser.add_argument("--width", type=int, required=True)
args = parser.parse_args()
H = args.height
W = args.width
MODEL = args.model
model = onnx.load(f'{MODEL}')
onnx_session = onnxruntime.InferenceSession(f'{MODEL}')
inputs = {}
for i in onnx_session.get_inputs():
inputs[i.name] = [i.shape[0], i.shape[1], H, W]
print('@@@@@@@@@@@@@@@@@@@@@ inputs')
pprint(inputs)
model_simp, check = simplify(model, input_shapes=inputs)
basename_without_ext = \
os.path.splitext(os.path.basename(MODEL))[0].split('_')[0]
onnx.save(model_simp, f'{basename_without_ext}_{H}x{W}.onnx')
import lbann.onnx.l2o
def parseInputShape(s):
name, shape = re.compile("^([^=]+)=([0-9,]+)$").search(s).groups()
return (name, list(map(int, shape.split(","))))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Convert a LBANN model to an ONNX model",
epilog="Usage: lbann2onnx.py model_alexnet.prototext output.onnx image=3,224,224 label=1000")
parser.add_argument("lbann_path", type=str,
help="Path to a LBANN model in .prototext")
parser.add_argument("onnx_path", type=str,
help="Path to an ONNX model")
parser.add_argument("input_shape", type=str, nargs="*",
help="Shape(s) of input tensor(s) *without* the mini-batch size in the \"NAME=N1,...,ND\" format.")
parser.add_argument("--add-value-info", dest="add_value_info", action="store_const",
const=True, default=False,
help="Embed value_info in the generated ONNX model")
args = parser.parse_args()
lbannPath = args.lbann_path
onnxPath = args.onnx_path
inputShapes = dict(map(parseInputShape, args.input_shape))
addValueInfo = args.add_value_info
model, miniBatchSize = lbann.onnx.l2o.parseLbannModelPB(os.path.expanduser(lbannPath),
inputShapes,
addValueInfo=addValueInfo)
onnx.save(model, os.path.expanduser(onnxPath))
[ # inputs
helper.make_tensor_value_info(
'input_1', TensorProto.FLOAT,
[batch_size, sequence_length, input_hidden_size]),
helper.make_tensor_value_info(
'input_2', TensorProto.FLOAT,
[batch_size, sequence_length, input_hidden_size]),
helper.make_tensor_value_info(
'input_mask',
TensorProto.FLOAT if use_float_mask else TensorProto.INT64,
[batch_size, sequence_length])
],
[ # outputs
helper.make_tensor_value_info(
'output', TensorProto.FLOAT,
[batch_size, sequence_length, input_hidden_size]),
],
initializers)
model = helper.make_model(graph)
return model
if __name__ == "__main__":
model = create_bert_attention()
onnx.save(model, "pruned_bert_attention.onnx")
model = create_bert_attention(switch_add_inputs=True)
onnx.save(model, "bert_attention_reverse_add_order.onnx")
model = create_tf2onnx_attention_3d()
onnx.save(model, "bert_3d_attention.onnx")
def generate_qat_support_model(model_names, test_initializers):
'''
EXPECTED_TEST_RESULT_CONFIG_1
'''
test_qat_support_models = []
# Main graph:
# [A] [input_bias]
# \ /
# Add [Transpose_output]
# \ |
# \ /
# Matmul -([input_weight])
# |
# |
# [B]
graph = helper.make_graph(
[ #nodes
helper.make_node("Add", ["A", "input_bias"], ["add_out"], "add0"),
helper.make_node("MatMul", ["add_out", "trans_out"], ["B"],
"matmul"),
],
"QAT_support_model_1", #name
[
#input
helper.make_tensor_value_info('A', TensorProto.FLOAT, ['unk_1'])
],
[
#output
helper.make_tensor_value_info('B', TensorProto.FLOAT, [1024])
])
#initializers
init_1 = test_initializers[0]
for init in init_1:
graph.initializer.add().CopyFrom(init)
model_1 = onnx.ModelProto()
model_1.ir_version = onnx.IR_VERSION
model_1 = onnx.helper.make_model(graph)
onnx.save(model_1, model_names[0])
test_qat_support_models.extend([model_1])
'''
EXPECTED_TEST_RESULT_CONFIG_2
'''
# Main graph:
# [A]
# |
# MaxPool
# / \
# Conv_0-[weight,bias] Conv_1-[weight,bias]
# \ /
# \ /
# Add
# |
# [B]
graph = helper.make_graph(
[ #nodes
helper.make_node("MaxPool", ["A"], ["maxpool_out"], "maxpool"),
helper.make_node("Conv", ["maxpool_out"], ["conv0_out"], "conv0"),
helper.make_node("Conv", ["maxpool_out"], ["conv1_out"], "conv1"),
helper.make_node("Add", ["conv0_out", "conv1_out"], ["B"], "add"),
],
"QAT_support_model_2", #name
[ #input
helper.make_tensor_value_info('A', TensorProto.FLOAT, ['unk_1'])
],
[ #output
helper.make_tensor_value_info('B', TensorProto.FLOAT,
[256, 64, 1, 1])
])
#initializers
init_2 = test_initializers[1]
for init in init_2:
graph.initializer.add().CopyFrom(init)
model_2 = onnx.ModelProto()
model_2.ir_version = onnx.IR_VERSION
model_2 = onnx.helper.make_model(graph)
onnx.save(model_1, model_names[1])
test_qat_support_models.extend([model_2])
return test_qat_support_models
def onnx_optimizer(onnx_net_path, input_shape=None):
onnx_net_opt_path = onnx_net_path[:-5] + '.opt.onnx'
print(os.getcwd())
print("----load onnx model: " + onnx_net_path)
onnx_model = onnx.load(onnx_net_path)
# all_passes = optimizer.get_available_passes()
# print("----available optimization passes:")
# for p in all_passes:
# print(p)
# print()
#
# print("----optimize onnx model: "+onnx_net_path)
# passes = ['eliminate_nop_pad',
# 'eliminate_identity',
# 'extract_constant_to_initializer',
# 'fuse_bn_into_conv',
# 'fuse_add_bias_into_conv',
# 'fuse_pad_into_conv',
# 'fuse_matmul_add_bias_into_gemm']
passes = [
'eliminate_identity', 'eliminate_nop_dropout',
'eliminate_nop_monotone_argmax', 'eliminate_nop_pad',
'eliminate_nop_transpose', 'extract_constant_to_initializer',
'fuse_bn_into_conv', 'fuse_add_bias_into_conv',
'fuse_consecutive_concats', 'fuse_consecutive_log_softmax',
'fuse_consecutive_reduce_unsqueeze', 'fuse_consecutive_squeezes',
'fuse_consecutive_transposes', 'fuse_matmul_add_bias_into_gemm',
'fuse_pad_into_conv', 'fuse_transpose_into_gemm'
]
#try:
# optimized_onnx_model = optimizer.optimize(onnx_model, passes)
# optimized_onnx_model = onnx.utils.polish_model(optimized_onnx_model)
#except IndexError as e:
# optimized_onnx_model = onnx_model
optimized_onnx_model = onnx_model
try:
input_shapes_ = {}
if input_shape is not None:
input_shape = input_shape.strip()
for x in input_shape.split(" "):
if ':' not in x:
input_shapes_[None] = list(map(int, x.split(',')))
else:
pieces = x.split(':')
# for the input name like input:0
name, shape = ':'.join(pieces[:-1]), list(
map(int, pieces[-1].split(',')))
input_shapes_[name] = shape
optimized_onnx_model, check_ok = onnx_simplifier.simplify(
optimized_onnx_model,
input_shapes=input_shapes_,
perform_optimization=False)
if not check_ok:
print("Check failed!")
exit()
except IndexError as e:
print(
"----onnxsim.simplify error: You'd better check the result with Netron"
)
print("----onnxsim.simplify error: " + str(RuntimeError))
except RuntimeError:
print(
"----onnxsim.simplify error: You'd better check the result with Netron"
)
print("----onnxsim.simplify error: " + str(RuntimeError))
# optimized_onnx_model = onnx_model
print("----export optimized onnx model: " + onnx_net_opt_path)
onnx.save(optimized_onnx_model, onnx_net_opt_path)
print("----export optimized onnx model done")
def keras2onnx2(all_file_onnx='mobilenet_save.onnx'):
onnx_model = onnxmltools.convert_keras(base_model)
onnx.save(onnx_model, all_file_onnx)
print('The model is invalid: %s' % e)
else:
print('The model is valid!')
if args.VERBOSE:
print('The model before conversion:\n{}'.format(original_model))
if args.no_convert:
quit()
# Opset version supported by current onnx-mlir
# Should be consistent with gen_onnx_mlir.py
current_onnx_mlir_support_version = 13
converted_model = version_converter.convert_version(
original_model, current_onnx_mlir_support_version)
if args.VERBOSE:
print('The model after conversion:\n{}'.format(converted_model))
if args.save:
inputFile = args.model
if inputFile.endswith('-opset' + str(current_onnx_mlir_support_version) +
'.onnx'):
printf('Converted model is not saved due to name conflict')
else:
outFile = inputFile[:inputFile.rfind(".onnx")] + '-opset-' + str(
current_onnx_mlir_support_version) + '.onnx'
onnx.save(converted_model, outFile)
print('The converted model is aved to ' + outFile)
gather_indice_initializer,
]
# Create the graph (GraphProto)
graph_def = helper.make_graph(
nodes,
"test-model",
[X, unsqueezed_masked_lm_positions],
[Y, Gather_Y],
initializers,
)
opsets = []
onnxdomain = OperatorSetIdProto()
onnxdomain.version = 12
onnxdomain.domain = "" # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
opsets.append(onnxdomain)
msdomain = OperatorSetIdProto()
msdomain.version = 1
msdomain.domain = "com.microsoft"
opsets.append(msdomain)
kwargs = {}
kwargs["opset_imports"] = opsets
model_def = helper.make_model(graph_def,
producer_name="onnx-example",
**kwargs)
onnx.save(model_def, "e2e.onnx")
assert len(inputs) == 1
input = inputs[0]
old_shape = []
new_shape = []
for i in range(len(input.type.tensor_type.shape.dim)):
nd = input.type.tensor_type.shape.dim[i].dim_value
old_shape.append(nd)
if i >= 2:
nd = int(nd * ratio)
new_shape.append(nd)
input.type.tensor_type.shape.dim[i].dim_value = nd
print('new_shape', new_shape)
onnx.save(model, os.path.join(output_dir, 'model.onnx'))
input_pb = os.path.join(output_dir, 'test_data_set_0/input_0.pb')
input_tensor = onnx.load_tensor(input_pb)
assert len(input_tensor.dims) == len(new_shape)
input = onnx.numpy_helper.to_array(input_tensor)
pad_width = []
for od, nd in zip(old_shape, new_shape):
pad_width.append((0, nd - od))
input = np.pad(input, pad_width, 'constant')
onnx.save_tensor(onnx.numpy_helper.from_array(input), input_pb)
def export_tensorflow(
tokenizer: PreTrainedTokenizer,
model: TFPreTrainedModel,
config: OnnxConfig,
opset: int,
output: Path,
) -> Tuple[List[str], List[str]]:
"""
Export a TensorFlow model to an ONNX Intermediate Representation (IR)
Args:
tokenizer ([`PreTrainedTokenizer`]):
The tokenizer used for encoding the data.
model ([`TFPreTrainedModel`]):
The model to export.
config ([`~onnx.config.OnnxConfig`]):
The ONNX configuration associated with the exported model.
opset (`int`):
The version of the ONNX operator set to use.
output (`Path`):
Directory to store the exported ONNX model.
Returns:
`Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
the ONNX configuration.
"""
import tensorflow as tf
import onnx
import tf2onnx
model.config.return_dict = True
# Check if we need to override certain configuration item
if config.values_override is not None:
logger.info(
f"Overriding {len(config.values_override)} configuration item(s)")
for override_config_key, override_config_value in config.values_override.items(
):
logger.info(
f"\t- {override_config_key} -> {override_config_value}")
setattr(model.config, override_config_key, override_config_value)
# Ensure inputs match
model_inputs = config.generate_dummy_inputs(
tokenizer, framework=TensorType.TENSORFLOW)
inputs_match, matched_inputs = ensure_model_and_config_inputs_match(
model, model_inputs.keys())
onnx_outputs = list(config.outputs.keys())
input_signature = [
tf.TensorSpec.from_tensor(tensor, name=key)
for key, tensor in model_inputs.items()
]
onnx_model, _ = tf2onnx.convert.from_keras(model,
input_signature,
opset=opset)
onnx.save(onnx_model, output.as_posix())
config.restore_ops()
return matched_inputs, onnx_outputs
# 删除输入
# for n in graph.input:
# if n.name == 'lr1':
# graph.input.remove(n)
# 指定node的输入输出
# node[77].input[0] = 'lr'
# node[47].output[0] = 'sr'
# node[50].output[0] = '29' # 节点序号
# 改变输入输出维度
# graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
# graph.input[0].type.tensor_type.shape.dim[1].dim_value = 48
# graph.input[0].type.tensor_type.shape.dim[2].dim_value = 180
# graph.input[0].type.tensor_type.shape.dim[3].dim_value = 320
# 新建节点，插入模型中
new_node = onnx.helper.make_node(
"ReshapeGuidemapAndGridSample",
inputs=['145', '140'],
outputs=['158'],
)
graph.node.insert(999, new_node)
# node[77].input = '158'
# 保存onnx
# onnx.checker.check_model(onnx_model) # 检查模型是否有异常，有时可以跳过
onnx.save(onnx_model, out_onnx_path)
msdomain = OperatorSetIdProto()
msdomain.version = 1
msdomain.domain = "com.microsoft"
opsets = [onnxdomain, msdomain]
onnx.save(
helper.make_model(
helper.make_graph(
[add, softmax_no_axis],
"Add_Softmax_Fusion",
[
helper.make_tensor_value_info("input", TensorProto.FLOAT,
["d_1", "d_2"]),
helper.make_tensor_value_info("bias", TensorProto.FLOAT,
["d_1", "d_2"]),
],
[
helper.make_tensor_value_info("output", TensorProto.FLOAT,
["d_1", "d_2"]),
],
[],
),
opset_imports=opsets,
),
r"bias_softmax_fusion_simple_no_axis_opset13.onnx",
)
onnx.save(
helper.make_model(
helper.make_graph(
[add, softmax1],
['A', 'B'], # inputs
['C'], # outputs,
)
# Create the graph (GraphProto)
graph_def = helper.make_graph(
[node_def],
'test-model',
[A, B],
[C],
)
# Create the model (ModelProto)
model_def = helper.make_model(graph_def, producer_name='onnx-example')
onnx.save(model_def, 'model.onnx')
##################################################
print('The model is:\n{}'.format(model_def))
onnx.checker.check_model(model_def)
print('The model is checked!')
##################################################
# create runtime session
sess = rt.InferenceSession("model.onnx")
print("-------------\n")
# get output name
input_name1 = sess.get_inputs()[0].name
import onnx
from onnx import helper, numpy_helper
import numpy as np
import copy
exported_model = onnx.load("source.onnx")
for node in exported_model.graph.node:
if node.op_type in ["RandomUniformLike"]:
attributes = list(node.attribute)
new_attributes = []
for attr in attributes:
if attr.name == "seed":
new_attributes.append(
helper.make_attribute('seed',
float(np.random.randint(1e6))))
else:
new_attributes.append(copy.deepcopy(attr))
del node.attribute[:]
node.attribute.extend(new_attributes)
onnx.save(exported_model, "modified.onnx")
import torch
import onnx
from onnx import shape_inference
from model import Discriminator
# MODEL_NAME = 'G_epoch_4_100_osrgan.pth'
# model = Generator_OSRGAN(4).eval()
# model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
# x = torch.randn(64, 3, 9, 9)
# file = 'epochs/' + 'ONNX_' + MODEL_NAME + '_OSRGAN.onnx'
# torch.onnx.export(model, x, file)
# onnx.save(onnx.shape_inference.infer_shapes(onnx.load(file)), file)
MODEL_NAME = 'D_epoch_4_100.pth'
model = Discriminator().eval()
model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
x = torch.randn(64, 3, 3, 3)
file = 'epochs/' + 'ONNX_' + MODEL_NAME + '.onnx'
torch.onnx.export(model, x, file)
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(file)), file)
DetectionCheckpointer(torch_model).resume_or_load(cfg.MODEL.WEIGHTS)
# get a sample data
data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
first_batch = next(iter(data_loader))
# convert and save caffe2 model
tracer = Caffe2Tracer(cfg, torch_model, first_batch)
if args.format == "caffe2":
caffe2_model = tracer.export_caffe2()
caffe2_model.save_protobuf(args.output)
# draw the caffe2 graph
caffe2_model.save_graph(os.path.join(args.output, "model.svg"), inputs=first_batch)
elif args.format == "onnx":
onnx_model = tracer.export_onnx()
onnx.save(onnx_model, os.path.join(args.output, "model.onnx"))
elif args.format == "torchscript":
script_model = tracer.export_torchscript()
script_model.save(os.path.join(args.output, "model.ts"))
# Recursively print IR of all modules
with open(os.path.join(args.output, "model_ts_IR.txt"), "w") as f:
try:
f.write(script_model._actual_script_module._c.dump_to_str(True, False, False))
except AttributeError:
pass
# Print IR of the entire graph (all submodules inlined)
with open(os.path.join(args.output, "model_ts_IR_inlined.txt"), "w") as f:
f.write(str(script_model.inlined_graph))
# Print the model structure in pytorch style
with open(os.path.join(args.output, "model.txt"), "w") as f:
def evaluate_model(checkpoint_path,
activation="elu",
gate_type_coarse="regular_conv",
gate_type_fine="regular_conv",
kernel_size=3):
# Load the checkpoint
folder = os.path.dirname(checkpoint_path)
basename = os.path.splitext(os.path.splitext(checkpoint_path)[0])[0]
netG = InpaintingGenerator(depth_factor=32,
inference=False,
activation=activation,
gate_type_coarse=gate_type_coarse,
gate_type_fine=gate_type_fine,
kernel_size=kernel_size)
netG_state_dict = torch.load(checkpoint_path)
netG.load_state_dict(netG_state_dict)
netG.to("cuda:0")
netG.eval()
# Load the test input and output
nn_img = torch.load('save/nn_img.pt') * 2. - 1.
nn_img = nn_img.to("cuda:0")
nn_msk = torch.load('save/nn_msk.pt')
nn_msk = nn_msk.to("cuda:0")
# Export to ONNX
onnx_path = os.path.join(folder, basename + ".onnx")
torch.onnx.export(netG, (nn_img, nn_msk),
onnx_path,
input_names=["img", "msk"],
output_names=["out", "fine", "coarse"],
opset_version=9)
# Simplify
onnx_model = onnx.load(onnx_path)
model_simple, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx_path = os.path.join(folder, basename + "_simple.onnx")
onnx.save(model_simple, onnx_path)
# Check the ONNX model
ort_session = onnxruntime.InferenceSession(
onnx_path, providers=["CUDAExecutionProvider"])
io_binding = ort_session.io_binding()
io_binding.bind_input(name='img',
device_type='cuda',
device_id=0,
element_type=np.float32,
shape=nn_img.shape,
buffer_ptr=nn_img.data_ptr())
io_binding.bind_input(name='msk',
device_type='cuda',
device_id=0,
element_type=np.float32,
shape=nn_msk.shape,
buffer_ptr=nn_img.data_ptr())
io_binding.bind_output('out')
io_binding.bind_output('fine')
io_binding.bind_output('coarse')
num_runs = 10
tic = time.time()
for i in range(num_runs):
# onnx_output = ort_session.run(None, ort_inputs)[0][0] * 0.5 + 0.5
ort_session.run_with_iobinding(io_binding)
toc = time.time() - tic
onnx_output = io_binding.copy_outputs_to_cpu()[0]
onnx_output = onnx_output[0] * 0.5 + 0.5
print("Inference time ONNX", toc / num_runs)
onnx_output = np.swapaxes(onnx_output, 0, 2)
onnx_output = np.swapaxes(onnx_output, 0, 1)
# Run Pytorch model
tic = time.time()
for i in range(num_runs):
pytorch_output, _, _ = netG(nn_img, nn_msk)
toc = time.time() - tic
print("Inference time Pytorch", toc / num_runs)
pytorch_output = pytorch_output[0].detach().cpu().numpy() * 0.5 + 0.5
pytorch_output = np.swapaxes(pytorch_output, 0, 2)
pytorch_output = np.swapaxes(pytorch_output, 0, 1)
for i in range(10):
arr1 = np.round(pytorch_output, decimals=i)
arr2 = np.round(onnx_output, decimals=i)
print("Comparison round", i, "decimals:", np.array_equal(arr1, arr2),
np.count_nonzero(np.sum(np.abs(arr1 - arr2), axis=2)))
pytorch_synthesis = np.zeros_like(pytorch_output)
pytorch_synthesis[nn_msk[0][0].detach().cpu().numpy(
) > 0] = pytorch_output[nn_msk[0][0].detach().cpu().numpy() > 0]
onnx_synthesis = np.zeros_like(onnx_output)
onnx_synthesis[nn_msk[0][0].detach().cpu().numpy() > 0] = onnx_output[
nn_msk[0][0].detach().cpu().numpy() > 0]
print("Mean absolute difference:",
np.mean(np.abs(pytorch_synthesis - onnx_synthesis)))
print(
"Mean relative difference:",
np.mean(
np.abs((pytorch_synthesis - onnx_synthesis) /
(pytorch_synthesis + epsilon))))
plt.figure()
ax = plt.subplot(121)
plt.imshow(pytorch_output)
plt.axis('off')
ax.set_title("Pytorch")
ax = plt.subplot(122)
plt.imshow(onnx_output)
plt.axis('off')
ax.set_title("ONNX")
plt.show()
return
y1 = conv21(y)
y2 = conv22(y)
y = autograd.cat((y1, y2), 1)
y = autograd.relu(y)
y = autograd.flatten(y)
y = linear(y)
loss = autograd.softmax_cross_entropy(y, t)
return loss, y
autograd.training = True
for epoch in range(epochs):
for i in range(batch_number):
inputs = tensor.Tensor(device=dev, data=x_train[
i * 100:(1 + i) * 100], stores_grad=False)
targets = tensor.Tensor(device=dev, data=y_train[
i * 100:(1 + i) * 100], requires_grad=False, stores_grad=False)
loss, y = forward(inputs, targets)
model = sonnx.get_onnx_model(loss, inputs, targets)
onnx.save(model, 'cnn.onnx')
accuracy_rate = accuracy(tensor.to_numpy(y),
tensor.to_numpy(targets))
if (i % 5 == 0):
print('accuracy is:', accuracy_rate, 'loss is:',
tensor.to_numpy(loss)[0])
for p, gp in autograd.backward(loss):
sgd.update(p, gp)
sgd.step()
break
def from_torch_module(
model: torch.nn.Module,
save_path: Path,
inputs: Iterable[IOShape],
outputs: Iterable[IOShape],
model_input: Optional[List] = None,
opset: int = 10,
optimize: bool = True,
override: bool = False,
):
"""Save a loaded model in ONNX.
Arguments:
model (nn.Module): PyTorch model.
save_path (Path): ONNX saved path.
inputs (Iterable[IOShape]): Model input shapes. Batch size is indicated at the dimension.
outputs (Iterable[IOShape]): Model output shapes.
model_input (Optional[List]) : Sample Model input data
TODO reuse inputs to pass model_input parameter later
opset (int): ONNX op set version.
optimize (bool): Flag to optimize ONNX network. Default to `True`.
override (bool): Flag to override if the file with path to `save_path` has existed. Default to `False`.
"""
if save_path.with_suffix('.onnx').exists():
if not override: # file exist yet override flag is not set
logger.info('Use cached model')
return True
export_kwargs = dict()
# assert batch size
batch_sizes = list(map(lambda x: x.shape[0], inputs))
if not all(batch_size == batch_sizes[0] for batch_size in batch_sizes):
raise ValueError('batch size for inputs (i.e. the first dimensions of `input.shape` are not consistent.')
batch_size = batch_sizes[0]
if batch_size == -1:
export_kwargs['dynamic_axes'] = {
'input': {0: 'batch_size'}, # variable length axes
'output': {0: 'batch_size'}
}
batch_size = 1
else:
assert batch_size > 0
model.eval()
save_path.parent.mkdir(parents=True, exist_ok=True)
save_path_with_ext = save_path.with_suffix('.onnx')
dummy_tensors, input_names, output_names = list(), list(), list()
for input_ in inputs:
dtype = model_data_type_to_torch(input_.dtype)
dummy_tensors.append(torch.rand(batch_size, *input_.shape[1:], requires_grad=True, dtype=dtype))
input_names.append(input_.name)
for output_ in outputs:
output_names.append(output_.name)
if model_input is None:
model_input = tuple(dummy_tensors)
try:
torch.onnx.export(
model, # model being run
model_input, # model input (or a tuple for multiple inputs)
save_path_with_ext, # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=opset, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names=input_names, # the model's input names
output_names=output_names, # the model's output names
keep_initializers_as_inputs=True,
**export_kwargs
)
if optimize:
onnx_model = onnx.load(str(save_path_with_ext))
network = ONNXConverter.optim_onnx(onnx_model)
onnx.save(network, str(save_path_with_ext))
logger.info('ONNX format converted successfully')
return True
except:
#TODO catch different types of error
logger.warning("This model is not supported as ONNX format")
return False
def export(inputs, y, file_path, model_name="sonnx"):
onnx_model = to_onnx_model(inputs, y, model_name)
onnx.save(onnx_model, file_path)
def chainer_to_onnx(chainer_model, onnx_export_path):
dummy_input = numpy.zeros((1, 3, chainer_model.img_size, chainer_model.img_size), dtype=numpy.float32)
onnx_chainer.export(chainer_model, dummy_input, filename=onnx_export_path)
onnx_model = onnx.load(onnx_export_path)
onnx_model_optimized = optimizer.optimize(onnx_model, ['fuse_bn_into_conv'])
onnx.save(onnx_model_optimized, onnx_export_path)
