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_batch_scan(self):
input_dim = 3
hidden_dim = 5
bidirectional = False
layers = 3
for onnx_opset_ver in [7, 13]:
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,
onnx_opset_ver=onnx_opset_ver,
)
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, providers=onnxrt.get_available_providers())
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, providers=onnxrt.get_available_providers())
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)
