def test_symbolic_shape_infer(self):
cwd = os.getcwd()
test_model_dir = os.path.join(cwd, '..', 'models')
for filename in Path(test_model_dir).rglob('*.onnx'):
if filename.name.startswith('.'):
continue # skip some bad model files
print("Running symbolic shape inference on : " + str(filename))
SymbolicShapeInference.infer_shapes(in_mp=onnx.load(str(filename)),
auto_merge=True,
int_max=100000,
guess_output_rank=True)
def test_mismatched_types(self):
graph = helper.make_graph(
[
helper.make_node(
"If", ["x"], ["out"],
name="if_node",
then_branch=helper.make_graph(
[
helper.make_node(
"Constant",
[],
["one_float"],
value=helper.make_tensor(
"one_float_value", TensorProto.FLOAT, [],
[1]),
)
],
"then",
[],
[
helper.make_tensor_value_info(
"one_float", TensorProto.FLOAT, [])
],
),
else_branch=helper.make_graph(
[
helper.make_node(
"Constant",
[],
["one_double"],
value=helper.make_tensor(
"one_double", TensorProto.DOUBLE, [], [1]),
)
],
"else",
[],
[
helper.make_tensor_value_info(
"one_double", TensorProto.DOUBLE, [])
],
))
],
"graph",
[helper.make_tensor_value_info("x", TensorProto.BOOL, [])],
[helper.make_tensor_value_info("out", TensorProto.FLOAT, [])],
)
model = helper.make_model(graph, producer_name="test_mismatched_types")
with self.assertRaisesRegex(ValueError, r"if_node.*FLOAT.*DOUBLE"):
SymbolicShapeInference.infer_shapes(model, auto_merge=True)
def _test_einsum_two_inputs_impl(self, input_0_shape, input_1_shape,
output_0_shape, eqn):
nodes = [
helper.make_node("Einsum", ["input_0", "input_1"], ["output_0"],
"einsum_0",
equation=eqn),
]
inputs = [
helper.make_tensor_value_info('input_0', TensorProto.FLOAT,
input_0_shape),
helper.make_tensor_value_info('input_1', TensorProto.FLOAT,
input_1_shape),
]
outputs = [
helper.make_tensor_value_info('output_0', TensorProto.FLOAT, None),
]
graph = helper.make_graph(nodes, "Einsum_Test", inputs, outputs, [])
model = helper.make_model(graph)
inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
expected_shapes = [
helper.make_tensor_value_info('output_0', TensorProto.FLOAT,
output_0_shape)
]
self._check_shapes(graph, inferred.graph, expected_shapes)
def test_softmax_cross_entropy_loss(self):
hidden_size = 1024
nodes = [
helper.make_node("SoftmaxCrossEntropyLoss",
inputs=["logits", "labels"],
outputs=["loss"]),
]
inputs = [
helper.make_tensor_value_info('logits', TensorProto.FLOAT,
['b', 's', hidden_size]),
helper.make_tensor_value_info('labels', TensorProto.INT32,
['b', 's']),
]
outputs = [
helper.make_tensor_value_info('loss', TensorProto.FLOAT, None),
]
graph = helper.make_graph(nodes, "SoftmaxCrossEntropyLoss_Test",
inputs, outputs, [])
model = helper.make_model(graph)
inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
expected_shapes = [
helper.make_tensor_value_info('loss', TensorProto.FLOAT, [])
]
self._check_shapes(graph, inferred.graph, expected_shapes)
def test_gather_indices(self):
graph = helper.make_graph(
[
helper.make_node(
"Constant",
[],
["data"],
"constant",
value=helper.make_tensor("input", TensorProto.FLOAT, [5], [0.0, 1.0, 2.0, 3.0, 4.0]),
),
helper.make_node("Gather", ["data", "indices"], ["output"], axis=0),
],
"Gather_Test",
[
helper.make_tensor_value_info("indices", TensorProto.INT64, ["b"]),
],
[
helper.make_tensor_value_info("output", TensorProto.FLOAT, ["b"]),
],
)
model = helper.make_model(graph, producer_name="Gather_Test_Model")
model.opset_import[0].version = 13
inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
expected_shapes = [
helper.make_tensor_value_info("data", TensorProto.FLOAT, [5]),
helper.make_tensor_value_info("output", TensorProto.FLOAT, ["b"]),
]
self._check_shapes(graph, inferred.graph, expected_shapes)
def test_unsqueeze_opset_13(self):
graph = helper.make_graph(
[
helper.make_node("Unsqueeze", ["input", "axes"], ["temp"]),
helper.make_node("Identity", ["temp"], ["output"]),
],
"Unsqueeze_Test",
[
helper.make_tensor_value_info("input", TensorProto.FLOAT, ["b", "s"]),
],
[
helper.make_tensor_value_info("output", TensorProto.FLOAT, ["b", "s", 1]),
],
[
helper.make_tensor("axes", TensorProto.INT64, [1], [-1]),
],
)
model = helper.make_model(graph, producer_name="Unsqueeze_Test_Model")
model.opset_import[0].version = 13
inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
expected_shapes = [
helper.make_tensor_value_info("temp", TensorProto.FLOAT, ["b", "s", 1]),
helper.make_tensor_value_info("output", TensorProto.FLOAT, ["b", "s", 1]),
]
self._check_shapes(graph, inferred.graph, expected_shapes)
def check_slice_of_concat(self, input_dims, start, end, step, expected_output_dim):
_dimstrmap = {dim: f"dim{i}" for i, dim in enumerate(input_dims)}
def dimstrmap(dim):
return _dimstrmap.get(dim, dim)
def get_initializer(name):
valuemap = {"zero": 0, "one": 1, "two": 2, "ten": 10, "intmax": 2**32}
value = -valuemap[name[4:]] if name.startswith("neg_") else valuemap[name]
return onnx.helper.make_tensor(name, TensorProto.INT64, [1], [value])
initializers = [
get_initializer(name)
for name in [
"zero",
"one",
"two",
"ten",
"intmax",
"neg_intmax",
"neg_one",
"neg_ten",
]
]
inputs = []
nodes = []
for i, dim in enumerate(input_dims):
inputs.append(onnx.helper.make_tensor_value_info(f"t{i}", TensorProto.FLOAT, ["B", dim]))
nodes.extend(
[
onnx.helper.make_node("Shape", [f"t{i}"], [f"shape{i}"]),
onnx.helper.make_node("Slice", [f"shape{i}", "one", "two", "zero", "one"], [f"dim{i}"]),
onnx.helper.make_node("Neg", [f"dim{i}"], [f"neg_dim{i}"]),
]
)
def make_concat_dims(concat_name, dims):
dims = [f"neg_{dimstrmap(dim[1:])}" if dim.startswith("-") else dimstrmap(dim) for dim in dims]
return onnx.helper.make_node("Concat", dims, [concat_name], axis=0)
nodes.extend(
[
onnx.helper.make_node("Concat", [inp.name for inp in inputs], ["concat"], axis=1),
make_concat_dims("starts", ["zero", start]),
make_concat_dims("ends", ["intmax", end]),
make_concat_dims("axes", ["zero", "one"]),
make_concat_dims("steps", ["one", step]),
onnx.helper.make_node("Slice", ["concat", "starts", "ends", "axes", "steps"], ["output"]),
]
)
output = onnx.helper.make_tensor_value_info("output", TensorProto.FLOAT, ["d1", "d2"])
graph_def = onnx.helper.make_graph(nodes, "graph", inputs, [output], initializer=initializers)
model = SymbolicShapeInference.infer_shapes(onnx.helper.make_model(graph_def))
output = unique_element(model.graph.output)
shape = [d.dim_param if d.dim_param else d.dim_value for d in output.type.tensor_type.shape.dim]
self.assertEqual(shape, ["B", expected_output_dim])
def test_embed_layer_norm(self):
hidden_size = 32
initializers = [
helper.make_tensor('word_embedding', TensorProto.FLOAT,
[100, hidden_size],
[1.0] * (100 * hidden_size)),
helper.make_tensor('position_embedding', TensorProto.FLOAT,
[20, hidden_size], [1.0] * (20 * hidden_size)),
helper.make_tensor('segment_embedding', TensorProto.FLOAT,
[2, hidden_size], [1.0] * (2 * hidden_size)),
helper.make_tensor('gamma', TensorProto.FLOAT, [hidden_size],
[1.0] * hidden_size),
helper.make_tensor('beta', TensorProto.FLOAT, [hidden_size],
[1.0] * hidden_size)
]
nodes = [
helper.make_node("EmbedLayerNormalization",
inputs=[
"input_ids", "segment_ids", "word_embedding",
"position_embedding", "segment_embedding",
"gamma", "beta"
],
outputs=["output", "mask_index"],
domain="com.microsoft"),
]
inputs = [
helper.make_tensor_value_info('input_ids', TensorProto.FLOAT,
['b', 's']),
helper.make_tensor_value_info('segment_ids', TensorProto.FLOAT,
['b', 's']),
]
outputs = [
helper.make_tensor_value_info('output', TensorProto.FLOAT, None),
helper.make_tensor_value_info('mask_index', TensorProto.INT32,
None),
]
graph = helper.make_graph(nodes, "Unsqueeze_Test", inputs, outputs,
initializers)
model = helper.make_model(graph)
inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
expected_shapes = [
helper.make_tensor_value_info('output', TensorProto.FLOAT,
['b', 's', hidden_size]),
helper.make_tensor_value_info('mask_index', TensorProto.INT32,
['b'])
]
self._check_shapes(graph, inferred.graph, expected_shapes)
def test_unsqueeze_opset_11(self):
graph = helper.make_graph([
helper.make_node("Unsqueeze", ["input"], ["temp"], axes=[0]),
helper.make_node("Identity", ["temp"], ["output"]),
], "Unsqueeze_Test", [
helper.make_tensor_value_info('input', TensorProto.FLOAT, ['b', 's']),
], [
helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 'b', 's']),
])
model = helper.make_model(graph, producer_name='Unsqueeze_Test_Model')
model.opset_import[0].version = 11
inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
expected_shapes = [
helper.make_tensor_value_info('temp', TensorProto.FLOAT, [1, 'b', 's']),
helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 'b', 's'])
]
self._check_shapes(graph, inferred.graph, expected_shapes)
onnx.save(mp, output_model)
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--mode',
help='The modification mode',
choices=['to_scan', 'remove_initializers_from_inputs'])
parser.add_argument('--input', help='The input model file', default=None)
parser.add_argument('--output', help='The output model file', default=None)
return parser.parse_args()
if __name__ == '__main__':
args = parse_arguments()
print('input model: ' + args.input)
print('output model ' + args.output)
if args.mode == 'to_scan':
print('Convert LSTM to Scan...')
convert_to_scan_model(args.input, args.output)
elif args.mode == 'remove_initializers_from_inputs':
print(
'Remove all initializers from input for model with IR version >= 4...'
)
remove_initializers_from_inputs(args.input, args.output)
else:
raise NotImplementedError('Unknown mode')
print('Running symbolic shape inference on output model')
SymbolicShapeInference.infer_shapes(args.output, args.output)
print('Done!')
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--mode',
help='The modification mode',
choices=['to_scan', 'remove_initializers_from_inputs'])
parser.add_argument('--input', help='The input model file', default=None)
parser.add_argument('--output', help='The output model file', default=None)
return parser.parse_args()
if __name__ == '__main__':
args = parse_arguments()
print('input model: ' + args.input)
print('output model ' + args.output)
if args.mode == 'to_scan':
print('Convert LSTM/GRU/RNN to Scan...')
convert_to_scan_model(args.input, args.output)
elif args.mode == 'remove_initializers_from_inputs':
print(
'Remove all initializers from input for model with IR version >= 4...'
)
remove_initializers_from_inputs(args.input, args.output)
else:
raise NotImplementedError('Unknown mode')
print('Running symbolic shape inference on output model')
SymbolicShapeInference.infer_shapes(args.output,
args.output,
auto_merge=True)
print('Done!')
def perf_test(rnn_type,
num_threads,
input_dim,
hidden_dim,
bidirectional,
layers,
seq_len,
batch_size,
top_n=5,
min_duration_seconds=10):
set_num_threads(num_threads)
model_name = '{}_i{}_h{}_{}_l{}_{}.onnx'.format(
rnn_type, input_dim, hidden_dim, 'bi' if bidirectional else '', layers,
'batched' if batch_size > 1 else 'no_batch')
generate_model(rnn_type, input_dim, hidden_dim, bidirectional, layers,
model_name, batch_size == 1)
feeds = {
'input': np.random.rand(seq_len, batch_size,
input_dim).astype(np.float32)
}
# run original model
sess = onnxruntime.InferenceSession(model_name)
count, duration, per_iter_cost = perf_run(
sess,
feeds,
min_counts=top_n,
min_duration_seconds=min_duration_seconds)
avg_rnn = top_n_avg(per_iter_cost, top_n)
print('perf_rnn {}: run for {} iterations, top {} avg {} ms'.format(
model_name, count, top_n, avg_rnn))
# run Scan model converted from original
from model_editor import convert_to_scan_model
from symbolic_shape_infer import SymbolicShapeInference
scan_model_name = os.path.splitext(model_name)[0] + '_scan.onnx'
convert_to_scan_model(model_name, scan_model_name)
# note that symbolic shape inference is needed because model has symbolic batch dim, thus init_state is ConstantOfShape
SymbolicShapeInference.infer_shapes(scan_model_name, scan_model_name)
sess = onnxruntime.InferenceSession(scan_model_name)
count, duration, per_iter_cost = perf_run(
sess,
feeds,
min_counts=top_n,
min_duration_seconds=min_duration_seconds)
avg_scan = top_n_avg(per_iter_cost, top_n)
print('perf_scan {}: run for {} iterations, top {} avg {} ms'.format(
scan_model_name, count, top_n, avg_scan))
# quantize Scan model to int8
from model_quantizer import convert_matmul_model
int8_model_name = os.path.splitext(model_name)[0] + '_int8.onnx'
convert_matmul_model(scan_model_name, int8_model_name)
SymbolicShapeInference.infer_shapes(int8_model_name, int8_model_name)
sess = onnxruntime.InferenceSession(int8_model_name)
count, duration, per_iter_cost = perf_run(
sess,
feeds,
min_counts=top_n,
min_duration_seconds=min_duration_seconds)
avg_int8 = top_n_avg(per_iter_cost, top_n)
print('perf_int8 {}: run for {} iterations, top {} avg {} ms'.format(
int8_model_name, count, top_n, avg_int8))
return avg_rnn, avg_scan, avg_int8