onnxruntime/onnxruntime/test/python/onnxruntime_test_python_nuphar.py
liqunfu addb4b8c2b
Liqun/speech model loop to scan (#6070)
Provide a tool to convert Loop to Scan for Nuphar performance
Fix Nuphar CI pipeline failures.

Co-authored-by: liqun <liqun@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
2021-01-05 15:15:23 -08:00

717 lines
32 KiB
Python

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# -*- coding: UTF-8 -*-
import numpy as np
import onnx
from onnx import helper, numpy_helper
import onnxruntime as onnxrt
from helper import get_name
import os
from onnxruntime.nuphar.rnn_benchmark import perf_test, generate_model
from onnxruntime.nuphar.model_tools import validate_with_ort
import shutil
import sys
import subprocess
import tarfile
import unittest
import urllib.request
def reference_gemm(a, b, c, alpha, beta, transA, transB):
a = a if transA == 0 else a.T
b = b if transB == 0 else b.T
y = alpha * np.dot(a, b) + beta * c
return y
def set_gemm_node_attrs(attrs, config):
if config['alpha'] != 1.0:
attrs['alpha'] = config['alpha']
if config['beta'] != 1.0:
attrs['beta'] = config['beta']
if config['transA']:
attrs['transA'] = 1
if config['transB']:
attrs['transB'] = 1
def generate_gemm_inputs_initializers(graph, config, added_inputs_initializers={}, extend=False):
M = config['M']
K = config['K']
N = config['N']
shape_a = [K, M] if config['transA'] else [M, K]
shape_b = [N, K] if config['transB'] else [K, N]
shape_c = [M, N]
# when A/B/C are graph input of the main graph which contains
# a Scan node, then they need an extra 'seq' dimension
input_shape_a = ['seq'] + shape_a if extend else shape_a
input_shape_b = ['seq'] + shape_b if extend else shape_b
input_shape_c = ['seq'] + shape_c if extend else shape_c
np.random.seed(12345)
a = np.random.ranf(shape_a).astype(np.float32)
b = np.random.ranf(shape_b).astype(np.float32)
c = np.random.ranf(shape_c).astype(np.float32) if config['withC'] else np.array(0)
init_a = a if config['initA'] else None
init_b = b if config['initB'] else None
init_c = c if config['initC'] else None
A = config['A']
B = config['B']
C = config['C']
# A is an initializer
if A in added_inputs_initializers:
a = added_inputs_initializers[A]
else:
added_inputs_initializers[A] = a
if init_a is not None:
graph.initializer.add().CopyFrom(numpy_helper.from_array(init_a, A))
else:
graph.input.add().CopyFrom(helper.make_tensor_value_info(A,
onnx.TensorProto.FLOAT,
input_shape_a))
# B is an initializer
if B in added_inputs_initializers:
b = added_inputs_initializers[B]
else:
added_inputs_initializers[B] = b
if init_b is not None:
graph.initializer.add().CopyFrom(numpy_helper.from_array(init_b, B))
else:
graph.input.add().CopyFrom(helper.make_tensor_value_info(B,
onnx.TensorProto.FLOAT,
input_shape_b))
if config['withC']:
if C in added_inputs_initializers:
c = added_inputs_initializers[C]
else:
added_inputs_initializers[C] = c
if init_c is not None:
graph.initializer.add().CopyFrom(numpy_helper.from_array(init_c, C))
else:
graph.input.add().CopyFrom(helper.make_tensor_value_info(C,
onnx.TensorProto.FLOAT,
input_shape_c))
return (a, b, c)
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 generate_gemm_node_subgraph(scan_body, scan_node_inputs, postfix, config, added_inputs):
M = config['M']
K = config['K']
N = config['N']
shape_a = [K, M] if config['transA'] else [M, K]
shape_b = [N, K] if config['transB'] else [K, N]
A = config['A']
B = config['B']
gemm_node_inputs = []
# A comes from the outer graph if it's an initializer
if config['initA']:
gemm_node_inputs.append(A)
else:
gemm_node_inputs.append(A + postfix)
if A not in added_inputs:
added_inputs[A] = 1
scan_node_inputs.append(A)
scan_body.input.add().CopyFrom(helper.make_tensor_value_info(A + postfix,
onnx.TensorProto.FLOAT,
shape_a))
# B comes from the outer graph if it's an initializer
if config['initB']:
gemm_node_inputs.append(B)
else:
gemm_node_inputs.append(B + postfix)
if B not in added_inputs:
added_inputs[B] = 1
scan_node_inputs.append(B)
scan_body.input.add().CopyFrom(helper.make_tensor_value_info(B + postfix,
onnx.TensorProto.FLOAT,
shape_b))
# C comes from Scan state
if config['withC']:
gemm_node_inputs.append('in_' + config['C'] + postfix)
attrs = {}
set_gemm_node_attrs(attrs, config)
node = helper.make_node('Gemm',
gemm_node_inputs,
[config['Y'] + postfix],
config['node_name'],
**attrs)
scan_body.node.add().CopyFrom(node)
def generate_gemm_scan_model(model_name, config1, config2):
model = onnx.ModelProto()
model.ir_version = 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 set_gemm_model_inputs(config, test_inputs, a, b, c):
if not config['initA']:
test_inputs[config['A']] = a
if not config['initB']:
test_inputs[config['B']] = b
if config['withC'] and not config['initC']:
test_inputs[config['C']] = c
class TestNuphar(unittest.TestCase):
def test_bidaf(self):
# download BiDAF model
cwd = os.getcwd()
bidaf_url = 'https://onnxzoo.blob.core.windows.net/models/opset_9/bidaf/bidaf.tar.gz'
cache_dir = os.path.join(os.path.expanduser("~"), '.cache', 'onnxruntime')
os.makedirs(cache_dir, exist_ok=True)
bidaf_local = os.path.join(cache_dir, 'bidaf.tar.gz')
if not os.path.exists(bidaf_local):
urllib.request.urlretrieve(bidaf_url, bidaf_local)
with tarfile.open(bidaf_local, 'r') as f:
f.extractall(cwd)
# verify accuracy of quantized model
bidaf_dir = os.path.join(cwd, 'bidaf')
bidaf_model = os.path.join(bidaf_dir, 'bidaf.onnx')
bidaf_scan_model = os.path.join(bidaf_dir, 'bidaf_scan.onnx')
bidaf_opt_scan_model = os.path.join(bidaf_dir, 'bidaf_opt_scan.onnx')
bidaf_int8_scan_only_model = os.path.join(bidaf_dir, 'bidaf_int8_scan_only.onnx')
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor', '--input', bidaf_model, '--output',
bidaf_scan_model, '--mode', 'to_scan'
],
check=True,
cwd=cwd)
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor', '--input', bidaf_scan_model, '--output',
bidaf_opt_scan_model, '--mode', 'opt_inproj'
],
check=True,
cwd=cwd)
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_quantizer', '--input', bidaf_opt_scan_model, '--output',
bidaf_int8_scan_only_model, '--only_for_scan'
],
check=True,
cwd=cwd)
# run onnx_test_runner to verify results
# use -M to disable memory pattern
onnx_test_runner = os.path.join(cwd, 'onnx_test_runner')
subprocess.run([onnx_test_runner, '-e', 'nuphar', '-M', '-c', '1', '-j', '1', '-n', 'bidaf', cwd], check=True, cwd=cwd)
# test AOT on the quantized model
if os.name not in ['nt', 'posix']:
return # don't run the rest of test if AOT is not supported
cache_dir = os.path.join(cwd, 'nuphar_cache')
if os.path.exists(cache_dir):
shutil.rmtree(cache_dir)
os.makedirs(cache_dir)
# prepare feed
feed = {}
for i in range(4):
tp = onnx.load_tensor(os.path.join(bidaf_dir, 'test_data_set_0', 'input_{}.pb'.format(i)))
feed[tp.name] = numpy_helper.to_array(tp)
for model in [bidaf_opt_scan_model, bidaf_int8_scan_only_model]:
nuphar_settings = 'nuphar_cache_path:{}'.format(cache_dir)
for isa in ['avx', 'avx2', 'avx512']:
onnxrt.capi._pybind_state.set_nuphar_settings(nuphar_settings + ', nuphar_codegen_target:' + isa)
sess = onnxrt.InferenceSession(model) # JIT cache happens when initializing session
cache_dir_content = os.listdir(cache_dir)
assert len(cache_dir_content) == 1
cache_versioned_dir = os.path.join(cache_dir, cache_dir_content[0])
so_name = os.path.basename(model) + '.so'
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.create_shared', '--input_dir', cache_versioned_dir,
'--output_name', so_name
],
check=True)
nuphar_settings = 'nuphar_cache_path:{}, nuphar_cache_so_name:{}, nuphar_cache_force_no_jit:{}'.format(
cache_dir, so_name, 'on')
onnxrt.capi._pybind_state.set_nuphar_settings(nuphar_settings)
sess = onnxrt.InferenceSession(model)
sess.run([], feed)
# test avx
nuphar_settings = 'nuphar_cache_path:{}, nuphar_cache_so_name:{}, nuphar_cache_force_no_jit:{}, nuphar_codegen_target:{}'.format(
cache_dir, so_name, 'on', 'avx')
onnxrt.capi._pybind_state.set_nuphar_settings(nuphar_settings)
sess = onnxrt.InferenceSession(model)
sess.run([], feed)
def test_bert_squad(self):
# download BERT_squad model
cwd = os.getcwd()
bert_squad_url = 'https://onnxzoo.blob.core.windows.net/models/opset_10/bert_squad/download_sample_10.tar.gz'
cache_dir = os.path.join(os.path.expanduser("~"), '.cache', 'onnxruntime')
os.makedirs(cache_dir, exist_ok=True)
bert_squad_local = os.path.join(cache_dir, 'bert_squad.tar.gz')
if not os.path.exists(bert_squad_local):
urllib.request.urlretrieve(bert_squad_url, bert_squad_local)
with tarfile.open(bert_squad_local, 'r') as f:
f.extractall(cwd)
# run symbolic shape inference on this model
# set int_max to 1,000,000 to simplify symbol computes for things like min(1000000, seq_len) -> seq_len
bert_squad_dir = os.path.join(cwd, 'download_sample_10')
bert_squad_model = os.path.join(bert_squad_dir, 'bertsquad10.onnx')
subprocess.run([
sys.executable, '-m', 'onnxruntime.tools.symbolic_shape_infer', '--input', bert_squad_model, '--output',
bert_squad_model, '--auto_merge', '--int_max=1000000'
],
check=True,
cwd=cwd)
# run onnx_test_runner to verify results
onnx_test_runner = os.path.join(cwd, 'onnx_test_runner')
subprocess.run([onnx_test_runner, '-e', 'nuphar', '-n', 'download_sample_10', cwd], check=True, cwd=cwd)
# run onnxruntime_perf_test, note that nuphar currently is not integrated with ORT thread pool, so set -x 1 to avoid thread confliction with OpenMP
onnxruntime_perf_test = os.path.join(cwd, 'onnxruntime_perf_test')
subprocess.run([onnxruntime_perf_test, '-e', 'nuphar', '-x', '1', '-t', '20', bert_squad_model, '1.txt'],
check=True,
cwd=cwd)
def test_rnn_benchmark(self):
# make sure benchmarking scripts works
# note: quantized model requires AVX2, otherwise it might be slow
avg_rnn, avg_scan, avg_int8 = perf_test('lstm',
num_threads=1,
input_dim=128,
hidden_dim=1024,
bidirectional=True,
layers=1,
seq_len=16,
batch_size=1,
min_duration_seconds=1)
avg_rnn, avg_scan, avg_int8 = perf_test('gru',
num_threads=1,
input_dim=128,
hidden_dim=1024,
bidirectional=False,
layers=2,
seq_len=16,
batch_size=3,
min_duration_seconds=1)
avg_rnn, avg_scan, avg_int8 = perf_test('rnn',
num_threads=1,
input_dim=128,
hidden_dim=1024,
bidirectional=False,
layers=3,
seq_len=16,
batch_size=2,
min_duration_seconds=1)
def test_batch_scan(self):
input_dim = 3
hidden_dim = 5
bidirectional = False
layers = 3
lstm_model_name = 'test_batch_rnn_lstm.onnx'
# create an LSTM model for generating baseline data
generate_model('lstm',
input_dim,
hidden_dim,
bidirectional,
layers,
lstm_model_name,
batch_one=False,
has_seq_len=True)
seq_len = 8
batch_size = 2
# prepare input
data_input = (np.random.rand(seq_len, batch_size, input_dim) * 2 - 1).astype(np.float32)
data_seq_len = np.random.randint(1, seq_len, size=(batch_size,), dtype=np.int32)
# run lstm as baseline
sess = onnxrt.InferenceSession(lstm_model_name)
first_lstm_data_output = sess.run([], {'input': data_input[:, 0:1, :], 'seq_len': data_seq_len[0:1]})
lstm_data_output = []
lstm_data_output = first_lstm_data_output
for b in range(1, batch_size):
lstm_data_output = lstm_data_output + sess.run([], {
'input': data_input[:, b:(b + 1), :],
'seq_len': data_seq_len[b:(b + 1)]
})
lstm_data_output = np.concatenate(lstm_data_output, axis=1)
# generate a batch scan model
scan_model_name = 'test_batch_rnn_scan.onnx'
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor', '--input', lstm_model_name, '--output',
scan_model_name, '--mode', 'to_scan'
],
check=True)
# run scan_batch with batch size 1
sess = onnxrt.InferenceSession(scan_model_name)
scan_batch_data_output = sess.run([], {'input': data_input[:, 0:1, :], 'seq_len': data_seq_len[0:1]})
assert np.allclose(first_lstm_data_output, scan_batch_data_output)
# run scan_batch with batch size 2
scan_batch_data_output = sess.run([], {'input': data_input, 'seq_len': data_seq_len})
assert np.allclose(lstm_data_output, scan_batch_data_output)
# run scan_batch with batch size 1 again
scan_batch_data_output = sess.run([], {'input': data_input[:, 0:1, :], 'seq_len': data_seq_len[0:1]})
assert np.allclose(first_lstm_data_output, scan_batch_data_output)
def test_gemm_to_matmul(self):
gemm_model_name_prefix = "gemm_model"
matmul_model_name_prefix = "matmul_model"
common_config = {
'node_name':'GemmNode',
'A':'inputA', 'B':'inputB', 'C':'inputC', 'Y':'output', 'M':2, 'K':3, 'N':4, 'withC':False,
'initA':False, 'initB':False, 'initC':False, 'alpha':1.0, 'beta':1.0, 'transA':0, 'transB':0
}
test_configs = [
{},
{'transA':1},
{'transB':1},
{'transA':1, 'transB':1},
{'withC':True},
{'withC':True, 'initC':True},
{'initA':True},
{'initB':True},
{'initA':True, 'initB':True},
{'initA':True, 'transA':1},
{'initB':True, 'transB':1},
{'initA':True, 'transA':1, 'initB':True, 'transB':1},
{'alpha':2.2},
{'transA':1, 'alpha':2.2},
{'initA':True, 'transA':1, 'alpha':2.2},
{'withC':True, 'beta':3.3},
{'withC':True, 'initC':True, 'beta':3.3},
{'initA':True, 'transA':1, 'alpha':2.2, 'withC':True, 'initC':True, 'beta':3.3},
{'transA':1, 'transB':1, 'alpha':2.2, 'withC':True, 'beta':3.3},
{'transA':1, 'transB':1, 'alpha':2.2, 'withC':True, 'initC':True, 'beta':3.3}
]
for i, config in enumerate(test_configs):
running_config = common_config.copy()
running_config.update(config)
gemm_model_name = gemm_model_name_prefix+str(i)+'.onnx'
matmul_model_name = matmul_model_name_prefix+str(i)+'.onnx'
a, b, c, expected_y = generate_gemm_model(gemm_model_name, running_config)
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor',
'--input', gemm_model_name,
'--output', matmul_model_name, '--mode', 'gemm_to_matmul'
], check=True)
sess = onnxrt.InferenceSession(matmul_model_name)
test_inputs = {}
set_gemm_model_inputs(running_config, test_inputs, a, b, c)
actual_y = sess.run([], test_inputs)
assert np.allclose(expected_y, actual_y)
def test_gemm_to_matmul_with_scan(self):
gemm_model_name_prefix = "gemm_scan_model"
matmul_model_name_prefix = "matmul_scan_model"
common_config = {
'M':2, 'K':3, 'N':4, 'withC':False, 'initA':False, 'initB':False, 'initC':False,
'alpha':1.0, 'beta':1.0, 'transA':0, 'transB':0
}
common_config1 = common_config.copy()
common_config1.update({
'node_name':'GemmNode1', 'A':'input1A', 'B':'input1B', 'C':'input1C', 'Y':'output1'
})
common_config2 = common_config.copy()
common_config2.update({
'node_name':'GemmNode2', 'A':'input2A', 'B':'input2B', 'C':'input2C', 'Y':'output2'
})
test_configs = [
({}, {}),
({'transA':1}, {'transB':1}),
({'transA':1, 'transB':1}, {'transA':1, 'transB':1}),
({'alpha':2.2}, {'alpha':3.3}),
({'transA':1, 'transB':1, 'alpha':2.2}, {'transA':1, 'transB':1, 'alpha':3.3}),
({'withC':True, 'initC':True}, {}),
({'withC':True, 'initC':True}, {'withC':True, 'initC':True}),
({'transA':1, 'transB':1, 'alpha':2.2, 'withC':True, 'initC':True, 'beta':1.2},
{'transA':1, 'transB':1, 'alpha':3.3, 'withC':True, 'initC':True, 'beta':4.1}),
({'initA':True}, {}),
({'initA':True}, {'initB':True}),
# FIXME: enable the test below after we fix some likely issue in graph partitioner
#({'initA':True, 'initB':True}, {}),
#({'initA':True, 'initB':True, 'transA':1}, {'initA':True, 'transB':1}),
#({'initA':True, 'transA':1, 'transB':1, 'alpha':2.2},
# {'initB':True, 'transA':1, 'transB':1, 'alpha':3.3}),
#({'initA':True, 'transA':1, 'transB':1, 'alpha':2.2, 'withC':True, 'initC':True},
# {'initB':True, 'transA':1, 'transB':1, 'alpha':3.3}),
#({'initA':True, 'transA':1, 'transB':1, 'alpha':2.2, 'withC':True, 'initC':True, 'beta':1.2},
# {'initB':True, 'transA':1, 'transB':1, 'alpha':3.3, 'withC':True, 'initC':True, 'beta':4.2}),
({'A':'inputA', 'initA':True}, {'A':'inputA', 'initA':True}),
({'B':'inputB', 'initB':True}, {'B':'inputB', 'initB':True}),
({'C':'inputC', 'withC':True, 'initC':True}, {'C':'inputC', 'withC':True, 'initC':True}),
({'transA':1, 'alpha':1.2, 'B':'inputB', 'initB':True, 'C':'inputC', 'withC':True, 'initC':True},
{'transA':1, 'alpha':2.2, 'B':'inputB', 'initB':True, 'C':'inputC', 'withC':True, 'initC':True}),
({'transB':1, 'alpha':1.2, 'B':'inputB'}, {'transB':1, 'alpha':2.2, 'B':'inputB'}),
({'transA':1, 'alpha':1.2, 'A':'inputA', 'B':'inputB'},
{'transA':1, 'alpha':2.2, 'A':'inputA', 'B':'inputB'}),
({'transA':1, 'alpha':1.2, 'A':'inputA', 'B':'inputB', 'C':'inputC1', 'withC':True, 'initC':True},
{'transA':1, 'alpha':2.2, 'A':'inputA', 'B':'inputB', 'C':'inputC2', 'withC':True, 'initC':True}),
]
for i, config in enumerate(test_configs):
config1, config2 = config
running_config1 = common_config1.copy()
running_config1.update(config1)
running_config2 = common_config2.copy()
running_config2.update(config2)
gemm_model_name = gemm_model_name_prefix+str(i)+'.onnx'
matmul_model_name = matmul_model_name_prefix+str(i)+'.onnx'
a1, b1, c1, a2, b2, c2 = generate_gemm_scan_model(gemm_model_name,
running_config1,
running_config2)
a1 = a1.reshape((1, ) + a1.shape)
b1 = b1.reshape((1, ) + b1.shape)
a2 = a2.reshape((1, ) + a2.shape)
b2 = b2.reshape((1, ) + b2.shape)
sess = onnxrt.InferenceSession(gemm_model_name)
test_inputs = {}
set_gemm_model_inputs(running_config1, test_inputs, a1, b1, c1)
set_gemm_model_inputs(running_config2, test_inputs, a2, b2, c2)
# run before model editing
expected_y = sess.run([], test_inputs)
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor',
'--input', gemm_model_name,
'--output', matmul_model_name, '--mode', 'gemm_to_matmul'
], check=True)
# run after model editing
sess = onnxrt.InferenceSession(matmul_model_name)
actual_y = sess.run([], test_inputs)
assert np.allclose(expected_y, actual_y, atol=1e-7)
print("finished " + matmul_model_name)
def test_loop_to_scan(self):
loop_model_filename = get_name("nuphar_tiny_model_with_loop_shape_infered.onnx")
scan_model_filename = "nuphar_tiny_model_with_loop_shape_infered_converted_to_scan.onnx"
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor',
'--input', loop_model_filename,
'--output', scan_model_filename, '--mode', 'loop_to_scan'
], check=True)
validate_with_ort(loop_model_filename, scan_model_filename)
def test_loop_to_scan_with_inconvertible_loop(self):
# nuphar_onnx_test_loop11_inconvertible_loop.onnx contains a Loop op with dynamic loop count.
# This Loop op cannot be converted to a Scan op.
# Set --keep_unconvertible_loop_ops option so conversion will not fail due to unconvertible loop ops.
loop_model_filename = get_name("nuphar_onnx_test_loop11_inconvertible_loop.onnx")
scan_model_filename = "nuphar_onnx_test_loop11_inconvertible_loop_unchanged.onnx"
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_editor',
'--input', loop_model_filename,
'--output', scan_model_filename, '--mode', 'loop_to_scan',
'--keep_unconvertible_loop_ops'
], check=True)
# onnxruntime is failing with:
# onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 :
# FAIL : Non-zero status code returned while running Loop node. Name:''
# Status Message: Inconsistent shape in loop output for output. Expected:{1} Got:{0}
# skip validate_with_ort for now
# validate_with_ort(loop_model_filename, scan_model_filename)
def test_loop_to_scan_tool(self):
loop_model_filename = get_name("nuphar_tiny_model_with_loop_shape_infered.onnx")
scan_model_filename = "nuphar_tiny_model_with_loop_shape_infered_converted_to_scan.onnx"
subprocess.run([
sys.executable, '-m', 'onnxruntime.nuphar.model_tools',
'--input', loop_model_filename,
'--output', scan_model_filename,
'--tool', 'convert_loop_to_scan_and_validate',
'--symbolic_dims', 'sequence=30'
], check=True)
validate_with_ort(loop_model_filename, scan_model_filename)
if __name__ == '__main__':
unittest.main()