def generate_model(rnn_type, input_dim, hidden_dim, bidirectional, layers, model_name, batch_one=True, has_seq_len=False, onnx_opset_ver=7):
model = onnx.ModelProto()
model.ir_version = IR_VERSION
opset = model.opset_import.add()
opset.domain == 'onnx'
opset.version = onnx_opset_ver
num_directions = 2 if bidirectional else 1
X = 'input'
model.graph.input.add().CopyFrom(helper.make_tensor_value_info(X, onnx.TensorProto.FLOAT, ['s', 1 if batch_one else 'b', input_dim]))
model.graph.initializer.add().CopyFrom(numpy_helper.from_array(np.asarray([0, 0, -1], dtype=np.int64), 'shape'))
if has_seq_len:
seq_len = 'seq_len'
model.graph.input.add().CopyFrom(helper.make_tensor_value_info(seq_len, onnx.TensorProto.INT32, [1 if batch_one else 'b',]))
gates = {'lstm':4, 'gru':3, 'rnn':1}[rnn_type]
for i in range(layers):
layer_input_dim = (input_dim if i == 0 else hidden_dim * num_directions)
model.graph.initializer.add().CopyFrom(numpy_helper.from_array(np.random.rand(num_directions, gates*hidden_dim, layer_input_dim).astype(np.float32), 'W'+str(i)))
model.graph.initializer.add().CopyFrom(numpy_helper.from_array(np.random.rand(num_directions, gates*hidden_dim, hidden_dim).astype(np.float32), 'R'+str(i)))
model.graph.initializer.add().CopyFrom(numpy_helper.from_array(np.random.rand(num_directions, 2*gates*hidden_dim).astype(np.float32), 'B'+str(i)))
layer_inputs = [X, 'W'+str(i), 'R'+str(i), 'B'+str(i)]
if has_seq_len:
layer_inputs += [seq_len]
layer_outputs = ['layer_output_'+str(i)]
model.graph.node.add().CopyFrom(helper.make_node(rnn_type.upper(), layer_inputs, layer_outputs, rnn_type+str(i), hidden_size=hidden_dim, direction='bidirectional' if bidirectional else 'forward'))
model.graph.node.add().CopyFrom(helper.make_node('Transpose', layer_outputs, ['transposed_output_'+str(i)], 'transpose'+str(i), perm=[0,2,1,3]))
model.graph.node.add().CopyFrom(helper.make_node('Reshape', ['transposed_output_'+str(i), 'shape'], ['reshaped_output_'+str(i)], 'reshape'+str(i)))
X = 'reshaped_output_'+str(i)
model.graph.output.add().CopyFrom(helper.make_tensor_value_info(X, onnx.TensorProto.FLOAT, ['s', 'b', hidden_dim * num_directions]))
model = shape_inference.infer_shapes(model)
onnx.save(model, model_name)
def convert_gemm_to_matmul(input_model, output_model):
in_mp = onnx.load(input_model)
out_mp = onnx.ModelProto()
out_mp.CopyFrom(in_mp)
out_mp.ir_version = 5 # update ir version to avoid requirement of initializer in graph input
out_mp.graph.ClearField('node')
nf = NodeFactory(out_mp.graph)
# gemm_to_matmul will generate transposed weights if the corresponding input
# comes from initializer. We keep a map between the original and converted
# ones in case the original initializer is shared between Gemm ops
converted_initializers = {}
for in_n in in_mp.graph.node:
if in_n.op_type == 'Gemm':
gemm_to_matmul(in_n, nf, converted_initializers)
continue
out_n = out_mp.graph.node.add()
out_n.CopyFrom(in_n)
if in_n.op_type == 'Scan' or in_n.op_type == 'Loop':
in_subgraph = NodeFactory.get_attribute(in_n, 'body')
out_subgraph = NodeFactory.get_attribute(out_n, 'body')
out_subgraph.ClearField('node')
scan_nf = NodeFactory(out_mp.graph, out_subgraph)
for in_sn in in_subgraph.node:
if in_sn.op_type == 'Gemm':
gemm_to_matmul(in_sn, scan_nf, converted_initializers)
continue
out_sn = out_subgraph.node.add()
out_sn.CopyFrom(in_sn)
onnx.save(out_mp, output_model)
def setUp(self):
self.run_id = 1
self.flow_id = 2
self.error_data = '{"%s": "%s"}' % (ERROR_KEY, ERROR_MESSAGE)
self.empty_data = '{}'
self.simple_data = '{"run_id": %s, "flow_id": %s}' % (self.run_id,
self.flow_id)
self.non_empty_graph_data = go.Scatter(x=[1],
y=[1],
mode='lines',
name='data')
self.none_data = None
self.non_empty_loading = ['load']
self.empty_loading = []
self.non_empty_error = ['error']
self.empty_error = []
self.display_hidden = 'none'
self.display_visible = ''
self.visible_style = {DISPLAY_KEY: ''}
self.hidden_style = {DISPLAY_KEY: 'none'}
self.onnx_model = onnx.ModelProto()
# Create the static folder if id does not exist
if not os.path.exists(STATIC_PATH):
os.mkdir(STATIC_PATH)
def update_flow_graph(n_clicks, flow_data_json, nr_loads):
if has_error_or_is_loading(n_clicks, flow_data_json, nr_loads):
return None, EMPTY_LOADED
flow_data = json.loads(flow_data_json)
# Define paths
dot_path = os.path.join(STATIC_PATH,
'graph_{}.dot'.format(flow_data[FLOW_ID_KEY]))
svg_path = os.path.join(STATIC_PATH,
'graph_{}.svg'.format(flow_data[FLOW_ID_KEY]))
# Recreate the passed onnx model from the dictionary
model = onnx.ModelProto()
json_format.ParseDict(flow_data[ONNX_MODEL_KEY], model)
pydot_graph = GetPydotGraph(
model.graph,
name=model.graph.name,
rankdir=".",
node_producer=GetOpNodeProducer(embed_docstring=True,
**GRAPH_EXPORT_STYLE))
# Simplify the generated network graph
pydot_graph = simplyfy_pydot_graph(pydot_graph)
# Export the graph to a dot file and use it to create an svg
pydot_graph.write_dot(dot_path)
os.system('dot -Tsvg {} -o {}'.format(dot_path, svg_path))
# Remove the unused dot file
os.remove(dot_path)
return html.Iframe(src='/static/graph_{}.svg'.format(flow_data[FLOW_ID_KEY]),
style={'width': '100%', 'height': '90vh'}), \
flow_data[FLOW_ID_KEY]
def load_graph(fname, target):
model_proto = onnx.ModelProto()
with open(fname, "rb") as f:
data = f.read()
model_proto.ParseFromString(data)
g = GraphUtil.create_graph_from_onnx_model(model_proto, target)
return g, model_proto
def generate_gemm_model(model_name, config):
model = onnx.ModelProto()
model.ir_version = 7 # use stable onnx ir version
opset = model.opset_import.add()
opset.version = 11
added_inputs_initializers = {}
(a, b, c) = generate_gemm_inputs_initializers(model.graph, config,
added_inputs_initializers)
node_inputs = [config["A"], config["B"]]
if config["withC"]:
node_inputs.append(config["C"])
attrs = {}
set_gemm_node_attrs(attrs, config)
node = helper.make_node("Gemm", node_inputs, [config["Y"]],
config["node_name"], **attrs)
model.graph.node.add().CopyFrom(node)
shape_output = [config["M"], config["N"]]
model.graph.output.add().CopyFrom(
helper.make_tensor_value_info(config["Y"], onnx.TensorProto.FLOAT,
shape_output))
# compute reference output
y = reference_gemm(a, b, c, config["alpha"], config["beta"],
config["transA"], config["transB"])
onnx.save(model, model_name)
return (a, b, c, y)
def Run(self, model_str: str, inputs_str: List[str]):
model = onnx.ModelProto()
model.ParseFromString(model_str)
def deserialize_tp(tp_str):
tp = onnx.TensorProto()
tp.ParseFromString(tp_str)
return tp
input_tps = map(deserialize_tp, inputs_str)
input_arrs = map(onnx.numpy_helper.to_array, input_tps)
input_names = [x.name for x in model.graph.input]
inputs = dict(zip(input_names, input_arrs))
sess_options = rt.SessionOptions()
sess_options.graph_optimization_level = rt.GraphOptimizationLevel(0)
sess_options.log_severity_level = 3
sess = rt.InferenceSession(
model.SerializeToString(),
sess_options=sess_options,
providers=["CPUExecutionProvider"],
)
output_names = [x.name for x in sess.get_outputs()]
run_options = rt.RunOptions()
run_options.log_severity_level = 3
output_arrs = sess.run(output_names, inputs, run_options=run_options)
return [
onnx.numpy_helper.from_array(x).SerializeToString()
for x in output_arrs
]
def generate_gemm_model(model_name, config):
model = onnx.ModelProto()
model.ir_version = onnx.IR_VERSION
opset = model.opset_import.add()
opset.version = 11
added_inputs_initializers = {}
(a, b, c) = generate_gemm_inputs_initializers(model.graph, config,
added_inputs_initializers)
node_inputs = [config['A'], config['B']]
if config['withC']:
node_inputs.append(config['C'])
attrs = {}
set_gemm_node_attrs(attrs, config)
node = helper.make_node('Gemm', node_inputs, [config['Y']],
config['node_name'], **attrs)
model.graph.node.add().CopyFrom(node)
shape_output = [config['M'], config['N']]
model.graph.output.add().CopyFrom(
helper.make_tensor_value_info(config['Y'], onnx.TensorProto.FLOAT,
shape_output))
# compute reference output
y = reference_gemm(a, b, c, config['alpha'], config['beta'],
config['transA'], config['transB'])
onnx.save(model, model_name)
return (a, b, c, y)
def Save(dir, func, feed, outputs):
if not os.path.exists(dir):
os.makedirs(dir)
onnx_file = os.path.join(dir, model_file)
func.save(onnx_file, C.ModelFormat.ONNX)
# onnx model may have different name for RNN initial states as inputs
cntk_to_actual_names = {}
with open(onnx_file, 'rb') as ff:
sf = ff.read()
model = onnx.ModelProto()
model.ParseFromString(sf)
for actual_input in model.graph.input:
actual_input_name = actual_input.name
for cntk_input in func.arguments:
cntk_name = cntk_input.uid
if actual_input_name.startswith(cntk_name):
cntk_to_actual_names[cntk_name] = actual_input_name
if type(feed) is not dict:
feed = {func.arguments[0]: feed}
if type(outputs) is not dict:
outputs = {func.outputs[0]: outputs}
test_data_dir = os.path.join(dir, data_dir)
if not os.path.exists(test_data_dir):
os.makedirs(test_data_dir)
SaveData(test_data_dir, 'input', func.arguments,
[feed[var] for var in func.arguments], cntk_to_actual_names)
SaveData(test_data_dir, 'output', func.outputs,
[outputs[var] for var in func.outputs])
def model_output_names(model_file_name):
with open(model_file_name, 'rb') as pfile:
data_str = pfile.read()
model_proto = onnx.ModelProto()
model_proto.ParseFromString(data_str)
output_names = [out.name for out in model_proto.graph.output]
return output_names
def save_model(proto, f, format=None, save_as_external_data=False, all_tensors_to_one_file=True, location=None, size_threshold=1024, convert_attribute=False):
if isinstance(proto, bytes):
proto = onnx._deserialize(proto, onnx.ModelProto())
if save_as_external_data:
convert_model_to_external_data(proto, all_tensors_to_one_file, location, size_threshold, convert_attribute)
s = onnx._serialize(proto)
onnx._save_bytes(s, f)
def export_and_recurse(node, attribute, output_dir, level):
name = node.name
name = name.replace('/', '_')
sub_model = onnx.ModelProto()
sub_model.graph.MergeFrom(attribute.g)
filename = 'L' + str(
level
) + '_' + node.op_type + '_' + attribute.name + '_' + name + '.onnx'
onnx.save_model(sub_model, os.path.join(output_dir, filename))
dump_subgraph(sub_model, output_dir, level + 1)
def export_and_recurse(node, attribute, output_dir, level):
name = node.name
name = name.replace("/", "_")
sub_model = onnx.ModelProto()
sub_model.graph.MergeFrom(attribute.g)
filename = "L" + str(
level
) + "_" + node.op_type + "_" + attribute.name + "_" + name + ".onnx"
onnx.save_model(sub_model, os.path.join(output_dir, filename))
dump_subgraph(sub_model, output_dir, level + 1)
def convert_matmul_model(input_model, output_model, only_for_scan=False, share_input_quantization=False, preset_str='asymm8_param0_input1', qcfg_json=None, export_qcfg_json=None):
preset_qcfgs = {'asymm8_param0_input1' : {'W' : dict(QuantizeConfig(signed=1, reserved_bits=0, type_bits=8)),
'X' : dict(QuantizeConfig(signed=0, reserved_bits=1, type_bits=8)),
'Symmetric' : 0},
'symm16_param3_input3' : {'W' : dict(QuantizeConfig(signed=1, reserved_bits=3, type_bits=16)),
'X' : dict(QuantizeConfig(signed=1, reserved_bits=3, type_bits=16)),
'Symmetric' : 1}}
default_qcfg = preset_qcfgs[preset_str]
in_mp = onnx.load(input_model)
qcfg_dict = {}
if qcfg_json and not export_qcfg_json:
with open(qcfg_json, 'r') as f:
qcfg_dict = json.load(f)
out_mp = onnx.ModelProto()
out_mp.CopyFrom(in_mp)
out_mp.ir_version = 5 # update ir version to avoid requirement of initializer in graph input
onnx_opset_ver = ensure_opset(out_mp, 10) # bump up to ONNX opset 10, which is required for MatMulInteger
ensure_opset(out_mp, 1, 'com.microsoft') # add MS domain for MatMulInteger16
out_mp.graph.ClearField('node')
nf = NodeFactory(out_mp.graph)
converted_weights = {} # remember MatMul weights that have been converted, in case of sharing
quantized_inputs = {} if share_input_quantization else None # remember quantized inputs that might be able to share between MatMuls
for in_n in in_mp.graph.node:
if upgrade_op(nf, in_n):
continue
if in_n.op_type == 'MatMul' and not only_for_scan:
if quantize_matmul_2d_with_weight(in_n, in_mp.graph, nf, converted_weights, quantized_inputs, qcfg_dict, export_qcfg_json, default_qcfg, onnx_opset_ver):
continue
out_n = out_mp.graph.node.add()
out_n.CopyFrom(in_n)
if in_n.op_type == 'Scan' or in_n.op_type == 'Loop':
in_subgraph = NodeFactory.get_attribute(in_n, 'body')
out_subgraph = NodeFactory.get_attribute(out_n, 'body')
out_subgraph.ClearField('node')
scan_nf = NodeFactory(out_mp.graph, out_subgraph)
subgraph_quantized_inputs = {} if share_input_quantization else None # remember quantized inputs that might be able to share between MatMuls
for in_sn in in_subgraph.node:
if in_sn.op_type == 'MatMul':
if quantize_matmul_2d_with_weight(in_sn, in_subgraph, scan_nf, converted_weights, subgraph_quantized_inputs, qcfg_dict, export_qcfg_json, default_qcfg, onnx_opset_ver):
continue
if upgrade_op(scan_nf, in_sn):
continue
out_sn = out_subgraph.node.add()
out_sn.CopyFrom(in_sn)
onnx.save(out_mp, output_model)
if export_qcfg_json:
with open(qcfg_json, 'w') as f:
f.write(json.dumps(qcfg_dict, indent=2))
def __init__(self,
model,
verbose=0,
int_max=2**31 - 1,
auto_merge=True,
guess_output_rank=False):
super().__init__(int_max, auto_merge, guess_output_rank, verbose)
self.model_ = onnx.ModelProto()
self.model_.CopyFrom(model)
self.all_shapes_inferred_ = False
self.inferred_ = False
def optimize(self):
# Remove raw initializer data
model = onnx.ModelProto()
model.ParseFromString(self.data)
for initializer in model.graph.initializer:
self.remove_tensor_data(initializer)
for node in model.graph.node:
for attribute in node.attribute:
if attribute.t:
self.remove_tensor_data(attribute.t)
self.data = model.SerializeToString()
def _preprocess(self, in_mp):
out_mp = onnx.ModelProto()
out_mp.CopyFrom(in_mp)
out_mp.graph.ClearField('node')
self.out_mp_ = out_mp
defined = set([
i.name
for i in list(in_mp.graph.input) + list(in_mp.graph.initializer)
])
pending_nodes = []
# returns True if no more ready nodes
def _insert_ready_nodes():
ready_nodes = [
pn for pn in pending_nodes
if all([i in defined for i in pn.input if i])
]
for rn in ready_nodes:
self.out_mp_.graph.node.add().CopyFrom(rn)
for o in rn.output:
defined.add(o)
pending_nodes.remove(rn)
return not ready_nodes
# constant op -> initializer, topological sort
for in_n in in_mp.graph.node:
if in_n.op_type == 'Constant':
t = get_attribute(in_n, 'value')
t.name = in_n.output[0]
self.out_mp_.graph.initializer.add().CopyFrom(t)
defined.add(t.name)
else:
pending_nodes.append(in_n)
_insert_ready_nodes()
while pending_nodes:
if _insert_ready_nodes():
break
if pending_nodes and self.verbose_ > 0:
print('SymbolicShapeInference: orphaned nodes discarded: ')
print(*[n.op_type + ': ' + n.output[0] for n in pending_nodes],
sep='\n')
self.initializers_ = dict([(i.name, i)
for i in self.out_mp_.graph.initializer])
self.known_vi_ = dict([(i.name, i)
for i in list(self.out_mp_.graph.input)])
self.known_vi_.update(
dict([(i.name,
helper.make_tensor_value_info(i.name, i.data_type,
list(i.dims)))
for i in self.out_mp_.graph.initializer]))
def build_model(graph, ir_version, opset_version):
# type: (Graph)->onnx.ModelProto
model_proto = onnx.ModelProto()
build_graph(graph, model_proto.graph)
model_proto.ir_version = ir_version
model_proto.opset_import.add()
model_proto.opset_import[0].version = opset_version
return model_proto
def model_parameter_names(model_file_name):
with open(model_file_name, 'rb') as pfile:
data_str = pfile.read()
model_proto = onnx.ModelProto()
model_proto.ParseFromString(data_str)
init_names = set([(i.name) for i in model_proto.graph.initializer])
param_names = [
input.name for input in model_proto.graph.input
if input.name not in init_names
]
return param_names
def GenScan():
np.random.seed(0)
feature = C.sequence.input_variable((3, ), np.float32)
model = C.layers.For(
range(4), lambda: C.layers.Recurrence(LSTM(2, use_scan=True)))(feature)
data_feature = np.random.rand(2, 5, 3).astype(np.float32)
data_output = np.asarray(model.eval(data_feature))
Save('test_Scan', model, data_feature, data_output)
# Currently CNTK only outputs batch == 1, do some editing
in_mp = onnx.load('test_Scan/model.onnx')
out_mp = onnx.ModelProto()
out_mp.CopyFrom(in_mp)
out_mp.graph.ClearField('initializer')
# change LSTM init_c/h into inputs to support truncated sequence
# as batch dimension is unknown on those data when building model
# note here we assume init_c/h starts from 0
# if not the case, user need to manually broadcast it for feed
num_inputs = 1
for i in in_mp.graph.initializer:
if i.name.startswith("Constant"):
shape = i.dims
shape[0] = 2
aa = np.zeros(shape, dtype=np.float32)
tp = numpy_helper.from_array(aa, i.name)
with open(
'test_Scan/test_data_set_0/input_' + str(num_inputs) +
'.pb', 'wb') as ff:
ff.write(tp.SerializeToString())
num_inputs = num_inputs + 1
else:
out_mp.graph.initializer.add().CopyFrom(i)
for vi in list(out_mp.graph.input) + list(out_mp.graph.output) + list(
out_mp.graph.value_info):
dim = vi.type.tensor_type.shape.dim
dim[len(dim) - 2].dim_param = 'batch'
for n in out_mp.graph.node:
if n.op_type == 'Scan':
body = [attr for attr in n.attribute if attr.name == 'body'][0]
for vi in list(body.g.input) + list(body.g.output) + list(
body.g.value_info):
dim = vi.type.tensor_type.shape.dim
dim[0].dim_param = 'batch'
onnx.save(out_mp, 'test_Scan/model.onnx', 'wb')
def stripModelFromPrivateData(m):
"""
Strips the given model from all private data and returns a copy.
This usually includes all tensors, i.e. all w and b.
"""
# Create new model and copy all from m
privatizedModel = onnx.ModelProto()
privatizedModel.CopyFrom(m)
# Clear the tensors from the model
del privatizedModel.graph.initializer[:]
# Return the privatized model
return privatizedModel
def load_graph(fname):
with open(fname, "rb") as f:
data = f.read()
model_proto = onnx.ModelProto()
model_proto.ParseFromString(data)
onnx_nodes = model_proto.graph.node
output_names = []
# some pytorch model had empty names - make one up
for node in onnx_nodes:
if not node.name:
node.name = tf2onnx.utils.make_name("was_empty")
g = Graph(onnx_nodes, output_shapes={}, dtypes={}, output_names=output_names)
for i in model_proto.graph.initializer:
v = numpy_helper.to_array(i)
name = i.name
g.initializers[name] = i
dtype = i.data_type
g.set_dtype(name, dtype)
g.set_shape(name, v.shape)
for i in model_proto.graph.input:
name = i.name
if name in g.initializers:
# ignore if it is not a model input
continue
shape = [j.dim_value if hasattr(i.type.tensor_type, "dim_value") else -1
for j in i.type.tensor_type.shape.dim]
dtype = i.type.tensor_type.elem_type
g.set_dtype(name, dtype)
g.set_shape(name, shape)
g.add_graph_input(name, dtype, shape)
for i in model_proto.graph.output:
name = i.name
shape = [j.dim_value if hasattr(i.type.tensor_type, "dim_value") else -1
for j in i.type.tensor_type.shape.dim]
dtype = i.type.tensor_type.elem_type
g.set_dtype(name, dtype)
g.set_shape(name, shape)
output_names.append(name)
# TODO: this is a hack in case a output name does not follow tensorflow convention
for node in g.get_nodes():
for name in node.output:
g._nodes_by_name[name] = node # pylint: disable=protected-access
return g, model_proto.producer_name
def _run(self, inputs_np):
inputs_np_dict = {k: v for k, v in inputs_np if k != ""}
model = onnx.ModelProto()
model.CopyFrom(omm.model)
sess_options = onnxruntime.SessionOptions()
session = onnxruntime.InferenceSession(model.SerializeToString(),
sess_options)
ort_outputs = session.run(None, inputs_np_dict)
model.graph.ClearField("value_info")
initializers = {i.name: i for i in model.graph.initializer}
for i in model.graph.input:
if i.name in initializers:
continue
for idx, d in enumerate(i.type.tensor_type.shape.dim):
if d.dim_param != "":
d.ClearField("dim_param")
d.dim_value = inputs_np_dict[i.name].shape[idx]
try:
model = SymbolicShapeInference.infer_shapes(
model, 2**31 - 1, True, True, 1)
except:
logging.warning("Shape infer by onnxruntime failed.")
with TemporaryDirectory() as tmpdir:
clear_op_code_generator()
model_code_generator = code_gen.get_model_code_generator(
model,
output_dir=tmpdir,
tensor_inplace=True,
simplify_names=True,
shape_infer=False)
model_code_generator.run()
spec = importlib.util.spec_from_file_location(
"model", os.path.join(tmpdir, "model.py"))
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
pt_outputs = mod.test_run_model(
[torch.from_numpy(v) for k, v in inputs_np if k != ""])
if type(pt_outputs) == torch.Tensor:
pt_outputs = [pt_outputs.detach().numpy()]
elif type(pt_outputs) in (list, tuple):
pt_outputs = [o.detach().numpy() for o in pt_outputs]
for l, r in zip(ort_outputs, pt_outputs):
assert np.allclose(l, r, atol=1e-4, rtol=1e-4, equal_nan=True)
def _run(self, inputs_np, onnx_model, gen_kwargs=None, tol=None):
inputs_np_dict = {k: v for k, v in inputs_np}
model = onnx.ModelProto()
model.CopyFrom(onnx_model)
sess_options = onnxruntime.SessionOptions()
session = onnxruntime.InferenceSession(model.SerializeToString(),
sess_options)
ort_outputs = session.run(None, inputs_np_dict)
model.graph.ClearField("value_info")
initializers = {i.name: i for i in model.graph.initializer}
for i in model.graph.input:
if i.name in initializers:
continue
for idx, d in enumerate(i.type.tensor_type.shape.dim):
if d.dim_param != "":
d.ClearField("dim_param")
d.dim_value = inputs_np_dict[i.name].shape[idx]
try:
model = SymbolicShapeInference.infer_shapes(model, 2**31 - 1, True, True,
1)
except:
logging.warning("Shape infer by onnxruntime failed.")
with TemporaryDirectory() as tmpdir:
if gen_kwargs is None:
gen_kwargs = {}
code_gen.gen(model,
output_dir=tmpdir,
tensor_inplace=False,
simplify_names=False,
shape_infer=False,
**gen_kwargs)
spec = importlib.util.spec_from_file_location(
"model", os.path.join(tmpdir, "model.py"))
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
pt_outputs = mod.test_run_model(
[torch.from_numpy(v) for _, v in inputs_np])
if tol is None:
tol = {"atol": 1e-5, "rtol": 1e-5}
for l, r in zip(ort_outputs, [o.detach().numpy() for o in pt_outputs]):
assert np.allclose(l, r, equal_nan=True, **tol)
def onnx_model(model=None):
"""Context manager that is the entry point to graph manipulations on model.
Manages the construction and destruction of the global model.
"""
global global_accessor
if global_accessor is not None:
raise RuntimeError(
"Base onnx model already exists. Cannot create multiple ModelAccessors."
)
# If the user did not provide a model, then assume that they want to build from scratch.
# It is the duty of the caller to fill the model however they deem fit.
if model is None:
model = onnx.ModelProto()
global_accessor = ModelAccessor(model)
try:
yield global_accessor
finally:
global_accessor = None
def _set_input_model(self, in_mp):
self._preprocess(in_mp)
self.initializers_ = dict([(i.name, i)
for i in self.out_mp_.graph.initializer])
self.known_vi_ = dict([(i.name, i)
for i in list(self.out_mp_.graph.input)])
self.known_vi_.update(
dict([(i.name,
helper.make_tensor_value_info(i.name, i.data_type,
list(i.dims)))
for i in self.out_mp_.graph.initializer]))
self.sympy_data_ = {}
self.dynamic_dims_ = {
} # new symbolic dims from some ops' dynamic output, i.e. NonZero
self.computed_dims_ = {}
# create a temporary ModelProto for single node inference
# note that we remove initializer to have faster inference
# for tensor ops like Reshape/Tile/Expand that read initializer, we need to do sympy computation based inference anyways
self.tmp_mp_ = onnx.ModelProto()
self.tmp_mp_.CopyFrom(self.out_mp_)
self.tmp_mp_.graph.ClearField('initializer')
def convert_to_scan_model(input_model, output_model):
in_mp = onnx.load(input_model)
out_mp = onnx.ModelProto()
out_mp.CopyFrom(in_mp)
out_mp.ir_version = 5 # update ir version to avoid requirement of initializer in graph input
ensure_opset(out_mp, 9) # bump up to ONNX opset 9, which is required for Scan
out_mp.graph.ClearField('node')
for in_n in in_mp.graph.node:
if in_n.op_type in ['LSTM', 'GRU', 'RNN']:
in_n = trim_unused_outputs(in_n, in_mp.graph)
if in_n.op_type == 'LSTM':
if convert_lstm_to_scan(in_n, out_mp.graph):
continue
if in_n.op_type == 'GRU':
if convert_gru_to_scan(in_n, out_mp.graph):
continue
if in_n.op_type == 'RNN':
if convert_rnn_to_scan(in_n, out_mp.graph):
continue
out_n = out_mp.graph.node.add()
out_n.CopyFrom(in_n)
onnx.save(out_mp, output_model)
def load_graph_def_from_pb(path):
with open(path, "rb") as f:
data = f.read()
model = onnx.ModelProto()
text_format.Parse(data, model)
return model.graph
def generate_model(
rnn_type,
input_dim,
hidden_dim,
bidirectional,
layers,
model_name,
batch_one=True,
has_seq_len=False,
onnx_opset_ver=7,
):
model = onnx.ModelProto()
model.ir_version = IR_VERSION
opset = model.opset_import.add()
opset.domain == "onnx"
opset.version = onnx_opset_ver
num_directions = 2 if bidirectional else 1
X = "input"
model.graph.input.add().CopyFrom(
helper.make_tensor_value_info(
X, onnx.TensorProto.FLOAT,
["s", 1 if batch_one else "b", input_dim]))
model.graph.initializer.add().CopyFrom(
numpy_helper.from_array(np.asarray([0, 0, -1], dtype=np.int64),
"shape"))
if has_seq_len:
seq_len = "seq_len"
model.graph.input.add().CopyFrom(
helper.make_tensor_value_info(
seq_len,
onnx.TensorProto.INT32,
[
1 if batch_one else "b",
],
))
gates = {"lstm": 4, "gru": 3, "rnn": 1}[rnn_type]
for i in range(layers):
layer_input_dim = input_dim if i == 0 else hidden_dim * num_directions
model.graph.initializer.add().CopyFrom(
numpy_helper.from_array(
np.random.rand(num_directions, gates * hidden_dim,
layer_input_dim).astype(np.float32),
"W" + str(i)))
model.graph.initializer.add().CopyFrom(
numpy_helper.from_array(
np.random.rand(num_directions, gates * hidden_dim,
hidden_dim).astype(np.float32), "R" + str(i)))
model.graph.initializer.add().CopyFrom(
numpy_helper.from_array(
np.random.rand(num_directions,
2 * gates * hidden_dim).astype(np.float32),
"B" + str(i)))
layer_inputs = [X, "W" + str(i), "R" + str(i), "B" + str(i)]
if has_seq_len:
layer_inputs += [seq_len]
layer_outputs = ["layer_output_" + str(i)]
model.graph.node.add().CopyFrom(
helper.make_node(
rnn_type.upper(),
layer_inputs,
layer_outputs,
rnn_type + str(i),
hidden_size=hidden_dim,
direction="bidirectional" if bidirectional else "forward",
))
model.graph.node.add().CopyFrom(
helper.make_node("Transpose",
layer_outputs, ["transposed_output_" + str(i)],
"transpose" + str(i),
perm=[0, 2, 1, 3]))
model.graph.node.add().CopyFrom(
helper.make_node("Reshape",
["transposed_output_" + str(i), "shape"],
["reshaped_output_" + str(i)],
"reshape" + str(i)))
X = "reshaped_output_" + str(i)
model.graph.output.add().CopyFrom(
helper.make_tensor_value_info(X, onnx.TensorProto.FLOAT,
["s", "b", hidden_dim * num_directions]))
model = shape_inference.infer_shapes(model)
onnx.save(model, model_name)
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, opset_imports=[helper.make_opsetid("", 13)])
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, opset_imports=[helper.make_opsetid("", 13)])
onnx.save(model_1, model_names[1])
test_qat_support_models.extend([model_2])
return test_qat_support_models
