mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-28 20:11:22 +00:00
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>
717 lines
32 KiB
Python
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()
|