From 56fd3225bbb7a5ec669fefc2d1922f3b1557ba8c Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Wed, 23 Jun 2021 13:49:53 -0700 Subject: [PATCH 01/27] modify lstm to be easily bidirectional --- python/tvm/relay/frontend/onnx.py | 134 ++++++++++++++++++------------ 1 file changed, 83 insertions(+), 51 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 5d07102f2c3f..80d10cd12379 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2057,6 +2057,52 @@ def _activation_needs_beta(cls, activation): class LSTM(RNN): """Operator converter for LSTM""" + @classmethod + def generate_lstm_forward(X_steps, H_t, C_t, W, R, B, p_i, p_f, p_o, f_act, g_act, h_act): + """Create an unrolled lstm loop. + + See https://github.com/onnx/onnx/blob/master/docs/Operators.md for math. + """ + h_list = [] + for step in X_steps: + step = _op.squeeze(step, axis=[0]) + gates = _op.nn.dense(step, W) + _op.nn.dense(H_t, R) + if B is not None: + WB, RB = _op.split(B, 2) + gates += WB + RB + i, o, f, c = _op.split(gates, 4, axis=-1) + + if p_i != 0: + i = f_act(i + p_i * C_t) + else: + i = f_act(i) + + if p_f != 0: + f = f_act(f + p_f * C_t) + else: + f = f_act(f) + + c = g_act(c) + C = f * C_t + i * c + if p_o != 0: + o = f_act(o + p_o * C) + else: + o = f_act(o) + + H = o * h_act(C) + + H_t = H + C_t = C + h_list.append(_op.expand_dims(H, axis=0)) + + # Concatenate outputs and add back in direction axis. + concatenated = _op.concatenate(h_list, 0) + output = _op.expand_dims(concatenated, axis=1) + H_t = _op.expand_dims(H_t, axis=0) + C_t = _op.expand_dims(C_t, axis=0) + + return output, H_t, C_t + @classmethod def _impl_v7(cls, inputs, attr, params): # Unpack inputs, note that if optional and not provided then value will be None. @@ -2066,20 +2112,15 @@ def _impl_v7(cls, inputs, attr, params): B = inputs[3] # Sequence length currently unused as it can be inferred from shapes. # sequence_lens = inputs['sequence_lens'] - h_0 = inputs[5] - c_0 = inputs[6] + H_0 = inputs[5] + C_0 = inputs[6] P = inputs[7] num_directions = infer_shape(W)[0] W_dtype = infer_type(W).checked_type.dtype - if num_directions != 1: + if num_directions not in [1, 2]: raise NotImplementedError("Bidirectional LSTMs not yet supported.") - # Remove num_directions axis from weights. - W = _op.squeeze(W, axis=[0]) - R = _op.squeeze(R, axis=[0]) - if B is not None: - B = _op.squeeze(B, axis=[0]) X_shape = infer_shape(X) hidden_size = infer_shape(R)[-1] @@ -2087,21 +2128,14 @@ def _impl_v7(cls, inputs, attr, params): # Initialize state if not provided. # Otherwise remove bidirectional axis. - if h_0 is None: - h_0 = _op.zeros((batch_size, hidden_size), W_dtype) - else: - h_0 = _op.squeeze(h_0, axis=[0]) - if c_0 is None: - c_0 = _op.zeros((batch_size, hidden_size), W_dtype) - else: - c_0 = _op.squeeze(c_0, axis=[0]) - + if H_0 is None: + H_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) + if C_0 is None: + C_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) if P is not None: - P = _op.squeeze(P, axis=[0]) - p_i, p_o, p_f = _op.split(P, 3) - H_t = h_0 - C_t = c_0 - h_list = [] + p_i, p_o, p_f = _op.split(P, 3, axis=1) + else: + p_i = p_o = p_f = _op.zeros((num_directions, hidden_size), W_dtype) if "activations" in attr: activations = attr["activations"] @@ -2134,37 +2168,35 @@ def _impl_v7(cls, inputs, attr, params): h_act = _op.tanh X_steps = _op.split(X, indices_or_sections=X_shape[0], axis=0) - for step in X_steps: - step = _op.squeeze(step, axis=[0]) - gates = _op.nn.dense(step, W) + _op.nn.dense(H_t, R) - if B is not None: - WB, RB = _op.split(B, 2) - gates += WB + RB - i, o, f, c = _op.split(gates, 4, axis=-1) - if P is not None: - i = f_act(i + p_i * C_t) - f = f_act(f + p_f * C_t) + result_output = [] + result_H = [] + result_C = [] + + for i in range(num_directions): + output, H, C = LSTM.generate_lstm_forward( + X=X_steps if i == 0 else X_steps[::-1], + H_t=H_0[i], + C_t=C_0[i], + W=W[i], + R=R[i], + B=B[i], + p_i=p_i[i], + p_f=p_f[i], + p_o=p_o[i], + f_act=f_act, + g_act=g_act, + h_act=h_act, + ) - else: - i = f_act(i) - f = f_act(f) - c = g_act(c) - C = f * C_t + i * c - if P is not None: - o = f_act(o + p_o * C) - else: - o = f_act(o) - H = o * h_act(C) - H_t = H - C_t = C - h_list.append(_op.expand_dims(H, axis=0)) - # Concatenate outputs and add back in direction axis. - concatenated = _op.concatenate(h_list, 0) - output = _op.expand_dims(concatenated, axis=1) - H_t = _op.expand_dims(H_t, axis=0) - C_t = _op.expand_dims(C_t, axis=0) + result_output.append(output) + result_H.append(H) + result_C.append(C) + + output = _op.concatenate(output, axis=1) + H = _op.concatenate(result_H, axis=0) + C = _op.concatenate(result_C, axis=0) - return _expr.TupleWrapper(_expr.Tuple((output, H_t, C_t)), 3) + return output, H, C class GRU(RNN): From 2447d35bfaceb2b78d69d9a11c273ff0b1c2e687 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 09:52:30 -0700 Subject: [PATCH 02/27] make it obvious some matriciies are packed via prime notation --- python/tvm/relay/frontend/onnx.py | 34 +++++++++++++++---------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 80d10cd12379..bf207a73ce58 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2107,20 +2107,20 @@ def generate_lstm_forward(X_steps, H_t, C_t, W, R, B, p_i, p_f, p_o, f_act, g_ac def _impl_v7(cls, inputs, attr, params): # Unpack inputs, note that if optional and not provided then value will be None. X = inputs[0] - W = inputs[1] - R = inputs[2] - B = inputs[3] + Wp = inputs[1] + Rp = inputs[2] + Bp = inputs[3] # Sequence length currently unused as it can be inferred from shapes. # sequence_lens = inputs['sequence_lens'] - H_0 = inputs[5] - C_0 = inputs[6] - P = inputs[7] + Hp_0 = inputs[5] + Cp_0 = inputs[6] + Pp = inputs[7] num_directions = infer_shape(W)[0] W_dtype = infer_type(W).checked_type.dtype if num_directions not in [1, 2]: - raise NotImplementedError("Bidirectional LSTMs not yet supported.") + raise ValueError("num_directions must be either 1 or 2!") X_shape = infer_shape(X) hidden_size = infer_shape(R)[-1] @@ -2128,11 +2128,11 @@ def _impl_v7(cls, inputs, attr, params): # Initialize state if not provided. # Otherwise remove bidirectional axis. - if H_0 is None: - H_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) - if C_0 is None: - C_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) - if P is not None: + if Hp_0 is None: + Hp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) + if Cp_0 is None: + Cp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) + if Pp is not None: p_i, p_o, p_f = _op.split(P, 3, axis=1) else: p_i = p_o = p_f = _op.zeros((num_directions, hidden_size), W_dtype) @@ -2175,11 +2175,11 @@ def _impl_v7(cls, inputs, attr, params): for i in range(num_directions): output, H, C = LSTM.generate_lstm_forward( X=X_steps if i == 0 else X_steps[::-1], - H_t=H_0[i], - C_t=C_0[i], - W=W[i], - R=R[i], - B=B[i], + H_t=Hp_0[i], + C_t=Cp_0[i], + W=Wp[i], + R=Rp[i], + B=Bp[i], p_i=p_i[i], p_f=p_f[i], p_o=p_o[i], From 85fd46d9de93e4a045537e94fb8348883f45929c Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 10:01:23 -0700 Subject: [PATCH 03/27] fix var name --- python/tvm/relay/frontend/onnx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index bf207a73ce58..9e4bf83f9b4a 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2116,8 +2116,8 @@ def _impl_v7(cls, inputs, attr, params): Cp_0 = inputs[6] Pp = inputs[7] - num_directions = infer_shape(W)[0] - W_dtype = infer_type(W).checked_type.dtype + num_directions = infer_shape(Wp)[0] + W_dtype = infer_type(Wp).checked_type.dtype if num_directions not in [1, 2]: raise ValueError("num_directions must be either 1 or 2!") From bcfb4d0e495d4aa7f73f0698858f5e0275926769 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 10:03:10 -0700 Subject: [PATCH 04/27] more var names --- python/tvm/relay/frontend/onnx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 9e4bf83f9b4a..cb0487412137 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2123,7 +2123,7 @@ def _impl_v7(cls, inputs, attr, params): raise ValueError("num_directions must be either 1 or 2!") X_shape = infer_shape(X) - hidden_size = infer_shape(R)[-1] + hidden_size = infer_shape(Rp)[-1] batch_size = X_shape[1] # Initialize state if not provided. @@ -2133,7 +2133,7 @@ def _impl_v7(cls, inputs, attr, params): if Cp_0 is None: Cp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) if Pp is not None: - p_i, p_o, p_f = _op.split(P, 3, axis=1) + p_i, p_o, p_f = _op.split(Pp, 3, axis=1) else: p_i = p_o = p_f = _op.zeros((num_directions, hidden_size), W_dtype) From d7133921dca16a4f6843479eb8229fccbeb6af68 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 10:42:30 -0700 Subject: [PATCH 05/27] add op split --- python/tvm/relay/frontend/onnx.py | 39 +++++++++++++++---------------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index cb0487412137..46ccb13e8182 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -34,18 +34,9 @@ from .. import qnn as _qnn from .. import ty as _ty from .. import vision as _vision -from .common import ( - AttrCvt, - Renamer, - fold_constant, - get_name, - get_relay_op, - infer_channels, - infer_shape, - infer_type, - infer_value, - new_var, -) +from .common import (AttrCvt, Renamer, fold_constant, get_name, get_relay_op, + infer_channels, infer_shape, infer_type, infer_value, + new_var) __all__ = ["from_onnx"] @@ -2172,17 +2163,25 @@ def _impl_v7(cls, inputs, attr, params): result_H = [] result_C = [] + H_ts = _op.split(Hp_0, num_directions) + C_ts = _op.split(Cp_0, num_directions) + Ws = _op.split(Wp, num_directions) + Rs = _op.split(Rp, num_directions) + Bs = _op.split(Bp, num_directions) + p_is = _op.split(p_i, num_directions) + p_fs = _op.split(p_f, num_directions) + p_os = _op.split(p_o, num_directions) for i in range(num_directions): output, H, C = LSTM.generate_lstm_forward( X=X_steps if i == 0 else X_steps[::-1], - H_t=Hp_0[i], - C_t=Cp_0[i], - W=Wp[i], - R=Rp[i], - B=Bp[i], - p_i=p_i[i], - p_f=p_f[i], - p_o=p_o[i], + H_t=H_ts[i], + C_t=C_ts[i], + W=Ws[i], + R=Rs[i], + B=Bs[i], + p_i=p_is[i], + p_f=p_fs[i], + p_o=p_os[i], f_act=f_act, g_act=g_act, h_act=h_act, From c8f61a407a746adc0501aeb42257d3fa03ec2daa Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 10:44:26 -0700 Subject: [PATCH 06/27] keyword arg names --- python/tvm/relay/frontend/onnx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 46ccb13e8182..9f9082063d6b 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2173,7 +2173,7 @@ def _impl_v7(cls, inputs, attr, params): p_os = _op.split(p_o, num_directions) for i in range(num_directions): output, H, C = LSTM.generate_lstm_forward( - X=X_steps if i == 0 else X_steps[::-1], + X_steps=X_steps if i == 0 else X_steps[::-1], H_t=H_ts[i], C_t=C_ts[i], W=Ws[i], From 26d1cd729f8e50e81dfc60efb0e6e2daa90350f7 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 10:46:09 -0700 Subject: [PATCH 07/27] missing implicit cls arg --- python/tvm/relay/frontend/onnx.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 9f9082063d6b..a7cda79450a2 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -34,9 +34,18 @@ from .. import qnn as _qnn from .. import ty as _ty from .. import vision as _vision -from .common import (AttrCvt, Renamer, fold_constant, get_name, get_relay_op, - infer_channels, infer_shape, infer_type, infer_value, - new_var) +from .common import ( + AttrCvt, + Renamer, + fold_constant, + get_name, + get_relay_op, + infer_channels, + infer_shape, + infer_type, + infer_value, + new_var, +) __all__ = ["from_onnx"] @@ -2049,7 +2058,7 @@ class LSTM(RNN): """Operator converter for LSTM""" @classmethod - def generate_lstm_forward(X_steps, H_t, C_t, W, R, B, p_i, p_f, p_o, f_act, g_act, h_act): + def generate_lstm_forward(cls, X_steps, H_t, C_t, W, R, B, p_i, p_f, p_o, f_act, g_act, h_act): """Create an unrolled lstm loop. See https://github.com/onnx/onnx/blob/master/docs/Operators.md for math. From 537727b7cff8e2f9b459c9222ac7c37d9f1095b5 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 13:45:16 -0700 Subject: [PATCH 08/27] deal with extra dimensions --- python/tvm/relay/frontend/onnx.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index a7cda79450a2..a9798bf956e5 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2181,16 +2181,24 @@ def _impl_v7(cls, inputs, attr, params): p_fs = _op.split(p_f, num_directions) p_os = _op.split(p_o, num_directions) for i in range(num_directions): + H_t = _op.squeeze(H_ts[i], axis=0) + C_t = _op.squeeze(C_ts[i], axis=0) + W = _op.squeeze(Ws[i], axis=0) + R = _op.squeeze(Rs[i], axis=0) + B = _op.squeeze(Bs[i], axis=0) + p_i = _op.squeeze(p_is[i], axis=0) + p_f = _op.squeeze(p_fs[i], axis=0) + p_o = _op.squeeze(p_os[i], axis=0) output, H, C = LSTM.generate_lstm_forward( X_steps=X_steps if i == 0 else X_steps[::-1], - H_t=H_ts[i], - C_t=C_ts[i], - W=Ws[i], - R=Rs[i], - B=Bs[i], - p_i=p_is[i], - p_f=p_fs[i], - p_o=p_os[i], + H_t=H_t, + C_t=C_t, + W=W, + R=R, + B=B, + p_i=p_i, + p_f=p_f, + p_o=p_o, f_act=f_act, g_act=g_act, h_act=h_act, @@ -2200,11 +2208,11 @@ def _impl_v7(cls, inputs, attr, params): result_H.append(H) result_C.append(C) - output = _op.concatenate(output, axis=1) + output = _op.concatenate(result_output, axis=1) H = _op.concatenate(result_H, axis=0) C = _op.concatenate(result_C, axis=0) - return output, H, C + return _expr.TupleWrapper(_expr.Tuple((output, H, C)), 3) class GRU(RNN): From ac3bb92e624152fc4f9808a39b6aa743247d274f Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 13:59:51 -0700 Subject: [PATCH 09/27] last of the fixes --- python/tvm/relay/frontend/onnx.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index a9798bf956e5..1bfdccea1ccf 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2132,6 +2132,8 @@ def _impl_v7(cls, inputs, attr, params): Hp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) if Cp_0 is None: Cp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) + if Bp is None: + Bp = _op.zeros((num_directions, hidden_size * 8), W_dtype) if Pp is not None: p_i, p_o, p_f = _op.split(Pp, 3, axis=1) else: @@ -2181,14 +2183,14 @@ def _impl_v7(cls, inputs, attr, params): p_fs = _op.split(p_f, num_directions) p_os = _op.split(p_o, num_directions) for i in range(num_directions): - H_t = _op.squeeze(H_ts[i], axis=0) - C_t = _op.squeeze(C_ts[i], axis=0) - W = _op.squeeze(Ws[i], axis=0) - R = _op.squeeze(Rs[i], axis=0) - B = _op.squeeze(Bs[i], axis=0) - p_i = _op.squeeze(p_is[i], axis=0) - p_f = _op.squeeze(p_fs[i], axis=0) - p_o = _op.squeeze(p_os[i], axis=0) + H_t = _op.squeeze(H_ts[i], axis=[0]) + C_t = _op.squeeze(C_ts[i], axis=[0]) + W = _op.squeeze(Ws[i], axis=[0]) + R = _op.squeeze(Rs[i], axis=[0]) + B = _op.squeeze(Bs[i], axis=[0]) + p_i = _op.squeeze(p_is[i], axis=[0]) + p_f = _op.squeeze(p_fs[i], axis=[0]) + p_o = _op.squeeze(p_os[i], axis=[0]) output, H, C = LSTM.generate_lstm_forward( X_steps=X_steps if i == 0 else X_steps[::-1], H_t=H_t, From a7ba617247f425cdcf0c986dc0aece4457d64e38 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 14:34:47 -0700 Subject: [PATCH 10/27] refactor rnn tests to support directions --- tests/python/frontend/onnx/test_forward.py | 152 ++++++++++++--------- 1 file changed, 85 insertions(+), 67 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 2f92f2d51994..ab87ec8f99c8 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import re + import numpy as np import pytest import scipy @@ -3175,88 +3176,105 @@ def verify_rnn( use_initial_state=False, use_peep=False, linear_before_reset=False, + directions=1, ): if rnn_type == "LSTM": multiplier = 4 elif rnn_type == "GRU": multiplier = 3 else: - raise NotImplementedError("%s RNNs not yet supported." % rnn_type) - x_np = np.random.uniform(size=(seq_length, batch_size, input_size)).astype("float32") - w_np = np.random.uniform(size=(1, multiplier * hidden_size, input_size)).astype("float32") - r_np = np.random.uniform(size=(1, multiplier * hidden_size, hidden_size)).astype("float32") - input_names = ["X", "W", "R"] - input_tensors = [ - helper.make_tensor_value_info("X", TensorProto.FLOAT, list(x_np.shape)), - helper.make_tensor_value_info("W", TensorProto.FLOAT, list(w_np.shape)), - helper.make_tensor_value_info("R", TensorProto.FLOAT, list(r_np.shape)), - ] - input_values = [x_np, w_np, r_np] - - if use_bias: - b_np = np.random.uniform(size=(1, multiplier * 2 * hidden_size)).astype("float32") - input_names.append("B") - input_tensors.append( - helper.make_tensor_value_info("B", TensorProto.FLOAT, [1, multiplier * 2 * hidden_size]) + raise NotImplementedError(f"{rnn_type} RNNs not yet supported.") + + if directions not in [1, 2]: + raise ValueError(f"Direction should be either 1 or 2 (for bidirectional LSTMs)") + + def get_inputs(): + input_names = [] + input_values = [] + input_tensors = [] + + def register(np_arr, name, shape=None): + input_values.append(np_arr) + input_names.append(name) + + # Map of numpy dtypes to the protobuf equivalent + dtype_map = { + "float32": TensorProto.FLOAT, + "int32": TensorProto.INT32, + "int8": TensorProto.INT8, + } + + if np_arr.dtype.name not in dtype_map: + raise ValueError(f"Unknown dtype we don't know how to handle {np.dtype.name}") + if shape is None: + shape = list(np_arr.shape) + proto_type = dtype_map[np_arr.dtype.name] + input_tensors.append(helper.make_tensor_value_info(name, proto_type, shape)) + + x_np = np.random.uniform(size=(seq_length, batch_size, input_size)).astype("float32") + w_np = np.random.uniform(size=(directions, multiplier * hidden_size, input_size)).astype( + "float32" ) - input_values.append(b_np) - - if use_initial_state: - assert use_bias == True, "Initial states must have bias specified." - sequence_np = np.repeat(seq_length, batch_size).astype("int32") - input_names.append("sequence_lens") - input_tensors.append( - helper.make_tensor_value_info("sequence_lens", TensorProto.INT32, [batch_size]) + r_np = np.random.uniform(size=(directions, multiplier * hidden_size, hidden_size)).astype( + "float32" ) - input_values.append(sequence_np) + register(x_np, "X") + register(w_np, "W") + register(r_np, "R") - initial_h_np = np.random.uniform(size=(1, batch_size, hidden_size)).astype("float32") - input_names.append("initial_h") - input_tensors.append( - helper.make_tensor_value_info( - "initial_h", TensorProto.FLOAT, [1, batch_size, hidden_size] + if use_bias: + b_np = np.random.uniform(size=(directions, multiplier * 2 * hidden_size)).astype( + "float32" ) - ) - input_values.append(initial_h_np) + register(b_np, "B") - if rnn_type == "LSTM": - initial_c_np = np.random.uniform(size=(1, batch_size, hidden_size)).astype("float32") - input_names.append("initial_c") - input_tensors.append( - helper.make_tensor_value_info( - "initial_c", TensorProto.FLOAT, [1, batch_size, hidden_size] - ) + if use_initial_state: + assert use_bias == True, "Initial states must have bias specified." + sequence_np = np.repeat(seq_length, batch_size).astype("int32") + register(sequence_np, "sequence_lens") + + initial_h_np = np.random.uniform(size=(directions, batch_size, hidden_size)).astype( + "float32" ) - input_values.append(initial_c_np) - - if use_peep and rnn_type == "LSTM": - assert use_initial_state == True, "Peepholes require initial state to be specified." - p_np = np.random.uniform(size=(1, 3 * hidden_size)).astype("float32") - input_names.append("P") - input_tensors.append( - helper.make_tensor_value_info("P", TensorProto.FLOAT, [1, 3 * hidden_size]) - ) - input_values.append(p_np) - - Y_shape = [seq_length, 1, batch_size, hidden_size] - Y_h_shape = [1, batch_size, hidden_size] - outputs = ["Y", "Y_h"] - graph_outputs = [ - helper.make_tensor_value_info("Y", TensorProto.FLOAT, list(Y_shape)), - helper.make_tensor_value_info("Y_h", TensorProto.FLOAT, list(Y_h_shape)), - ] - output_shapes = [Y_shape, Y_h_shape] + register(initial_h_np, "initial_h") - if rnn_type == "LSTM": - Y_c_shape = [1, batch_size, hidden_size] - outputs.append("Y_c") - graph_outputs.append( - helper.make_tensor_value_info("Y_c", TensorProto.FLOAT, list(Y_c_shape)) - ) - output_shapes.append(Y_c_shape) + if rnn_type == "LSTM": + initial_c_np = np.random.uniform(size=(directions, batch_size, hidden_size)).astype( + "float32" + ) + register(initial_c_np, "initial_c") + + if use_peep and rnn_type == "LSTM": + assert use_initial_state == True, "Peepholes require initial state to be specified." + p_np = np.random.uniform(size=(directions, 3 * hidden_size)).astype("float32") + register(p_np, "P") + + return input_names, input_tensors, input_values + + input_names, input_tensors, input_values = get_inputs() + + def get_outputs(): + output_names = [] + graph_outputs = [] + output_shapes = [] + + def register(name, shape, proto_type): + output_names.append(name) + graph_outputs.append(helper.make_tensor_value_info(name, proto_type, list(shape))) + output_shapes.append(list(shape)) + + register("Y", [seq_length, directions, batch_size, hidden_size], TensorProto.FLOAT) + register("Y_h", [directions, batch_size, hidden_size], TensorProto.FLOAT) + + if rnn_type == "LSTM": + register("Y_c", [directions, batch_size, hidden_size], TensorProto.FLOAT) + + return output_names, graph_outputs, output_shapes + + output_names, graph_outputs, output_shapes = get_outputs() rnn_node = helper.make_node( - rnn_type, inputs=input_names, outputs=outputs, hidden_size=hidden_size + rnn_type, inputs=input_names, outputs=output_names, hidden_size=hidden_size ) if activations is not None: activations_attr = helper.make_attribute("activations", activations) From a3dd53d3f26188e6044d374ff4bb8156722bf118 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 14:42:37 -0700 Subject: [PATCH 11/27] bidirectional tests --- tests/python/frontend/onnx/test_forward.py | 36 ++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index ab87ec8f99c8..cd68518bf1e9 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3381,6 +3381,42 @@ def test_lstm(): rnn_type="LSTM", ) + # Testing for bidirectionality + verify_rnn( + seq_length=3, + batch_size=3, + input_size=16, + hidden_size=40, + use_bias=True, + rnn_type="LSTM", + directions=2, + ) + + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "LeakyRelu", "Affine"], + alphas=[2.0, 0.5, 0.8], + betas=[0.3, 0.1], + rnn_type="LSTM", + directions=2, + ) + + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=True, + use_initial_state=True, + use_peep=True, + rnn_type="LSTM", + directions=2, + ) + @tvm.testing.uses_gpu def test_gru(): From 91c64833296e148330c08f0ca847acf3061b7571 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 14:46:44 -0700 Subject: [PATCH 12/27] test forward results --- tests/python/frontend/onnx/test_forward.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index cd68518bf1e9..c905ad60f60d 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3279,6 +3279,9 @@ def register(name, shape, proto_type): if activations is not None: activations_attr = helper.make_attribute("activations", activations) rnn_node.attribute.append(activations_attr) + if directions == 2: + direction_attr = helper.make_attribute("direction", "bidirectional") + rnn_node.attribute.append(direction_attr) if alphas is not None: alphas_attr = helper.make_attribute("activation_alpha", alphas) rnn_node.attribute.append(alphas_attr) From 1cdc2ae5b97d49dfe10f54b046f7945e25720601 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 14:55:59 -0700 Subject: [PATCH 13/27] go backwards --- python/tvm/relay/frontend/onnx.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 1bfdccea1ccf..98332ea3f68d 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2058,13 +2058,16 @@ class LSTM(RNN): """Operator converter for LSTM""" @classmethod - def generate_lstm_forward(cls, X_steps, H_t, C_t, W, R, B, p_i, p_f, p_o, f_act, g_act, h_act): + def generate_lstm( + cls, X_steps, H_t, C_t, W, R, B, p_i, p_f, p_o, f_act, g_act, h_act, backwards=False + ): """Create an unrolled lstm loop. See https://github.com/onnx/onnx/blob/master/docs/Operators.md for math. """ h_list = [] - for step in X_steps: + for i in len(X_steps): + step = X_steps[i] if not backwards else X_steps[-(i + 1)] step = _op.squeeze(step, axis=[0]) gates = _op.nn.dense(step, W) + _op.nn.dense(H_t, R) if B is not None: @@ -2191,8 +2194,8 @@ def _impl_v7(cls, inputs, attr, params): p_i = _op.squeeze(p_is[i], axis=[0]) p_f = _op.squeeze(p_fs[i], axis=[0]) p_o = _op.squeeze(p_os[i], axis=[0]) - output, H, C = LSTM.generate_lstm_forward( - X_steps=X_steps if i == 0 else X_steps[::-1], + output, H, C = LSTM.generate_lstm( + X_steps=X_steps, H_t=H_t, C_t=C_t, W=W, @@ -2204,6 +2207,7 @@ def _impl_v7(cls, inputs, attr, params): f_act=f_act, g_act=g_act, h_act=h_act, + backwards=i == 1, ) result_output.append(output) From 3382508210360137b17f64346658b187b6c8157e Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 15:08:52 -0700 Subject: [PATCH 14/27] more fixes --- python/tvm/relay/frontend/onnx.py | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 98332ea3f68d..ee2619b7527a 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -34,18 +34,9 @@ from .. import qnn as _qnn from .. import ty as _ty from .. import vision as _vision -from .common import ( - AttrCvt, - Renamer, - fold_constant, - get_name, - get_relay_op, - infer_channels, - infer_shape, - infer_type, - infer_value, - new_var, -) +from .common import (AttrCvt, Renamer, fold_constant, get_name, get_relay_op, + infer_channels, infer_shape, infer_type, infer_value, + new_var) __all__ = ["from_onnx"] @@ -2066,8 +2057,9 @@ def generate_lstm( See https://github.com/onnx/onnx/blob/master/docs/Operators.md for math. """ h_list = [] - for i in len(X_steps): - step = X_steps[i] if not backwards else X_steps[-(i + 1)] + seq_length = len(X_steps) + for i in range(seq_length): + step = X_steps[i] if not backwards else X_steps[seq_length-(i + 1)] step = _op.squeeze(step, axis=[0]) gates = _op.nn.dense(step, W) + _op.nn.dense(H_t, R) if B is not None: From ebc7078acf969e39ca60e67f1d78400c8d3fc10d Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 15:26:04 -0700 Subject: [PATCH 15/27] reverse tokens on reverse pass --- python/tvm/relay/frontend/onnx.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index ee2619b7527a..0308b91abe20 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -34,9 +34,18 @@ from .. import qnn as _qnn from .. import ty as _ty from .. import vision as _vision -from .common import (AttrCvt, Renamer, fold_constant, get_name, get_relay_op, - infer_channels, infer_shape, infer_type, infer_value, - new_var) +from .common import ( + AttrCvt, + Renamer, + fold_constant, + get_name, + get_relay_op, + infer_channels, + infer_shape, + infer_type, + infer_value, + new_var, +) __all__ = ["from_onnx"] @@ -2059,7 +2068,7 @@ def generate_lstm( h_list = [] seq_length = len(X_steps) for i in range(seq_length): - step = X_steps[i] if not backwards else X_steps[seq_length-(i + 1)] + step = X_steps[i] if not backwards else X_steps[seq_length - (i + 1)] step = _op.squeeze(step, axis=[0]) gates = _op.nn.dense(step, W) + _op.nn.dense(H_t, R) if B is not None: @@ -2090,6 +2099,10 @@ def generate_lstm( C_t = C h_list.append(_op.expand_dims(H, axis=0)) + if backwards: + # Canonical view is hidden states from the first token not last + h_list = h_list[::-1] + # Concatenate outputs and add back in direction axis. concatenated = _op.concatenate(h_list, 0) output = _op.expand_dims(concatenated, axis=1) From 023da4f54dc006c734f98b1cd88d701d57f5e330 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 15:29:43 -0700 Subject: [PATCH 16/27] parameterized directions --- tests/python/frontend/onnx/test_forward.py | 236 +++++++++++---------- 1 file changed, 121 insertions(+), 115 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index c905ad60f60d..5a6a7f8bb864 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3301,124 +3301,130 @@ def register(name, shape, proto_type): @tvm.testing.uses_gpu def test_lstm(): - # No bias. - verify_rnn( - seq_length=2, batch_size=1, input_size=16, hidden_size=32, use_bias=False, rnn_type="LSTM" - ) - # large batch. - verify_rnn( - seq_length=4, batch_size=8, input_size=16, hidden_size=32, use_bias=True, rnn_type="LSTM" - ) - # Non power of two. - verify_rnn( - seq_length=3, batch_size=3, input_size=16, hidden_size=40, use_bias=True, rnn_type="LSTM" - ) - # Long sequence. - verify_rnn( - seq_length=8, batch_size=1, input_size=16, hidden_size=32, use_bias=True, rnn_type="LSTM" - ) - # Large hidden. - verify_rnn( - seq_length=2, batch_size=1, input_size=16, hidden_size=128, use_bias=True, rnn_type="LSTM" - ) - # Large input. - verify_rnn( - seq_length=2, batch_size=1, input_size=64, hidden_size=32, use_bias=True, rnn_type="LSTM" - ) - - # Different activation testing. - # Default value hardsigmoid. - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "Tanh", "Tanh"], - rnn_type="LSTM", - ) - # Multiple parameterized activations. - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "LeakyRelu", "Tanh"], - alphas=[2.0, 0.5], - betas=[0.3], - rnn_type="LSTM", - ) - # All parameterized with new Affine activation. - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "LeakyRelu", "Affine"], - alphas=[2.0, 0.5, 0.8], - betas=[0.3, 0.1], - rnn_type="LSTM", - ) - - # Testing with initial state and peepholes - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=True, - use_initial_state=True, - rnn_type="LSTM", - ) - - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=True, - use_initial_state=True, - use_peep=True, - rnn_type="LSTM", - ) + for directions in [1, 2]: + # No bias. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + rnn_type="LSTM", + directions=directions, + ) + # large batch. + verify_rnn( + seq_length=4, + batch_size=8, + input_size=16, + hidden_size=32, + use_bias=True, + rnn_type="LSTM", + directions=directions, + ) + # Non power of two. + verify_rnn( + seq_length=3, + batch_size=3, + input_size=16, + hidden_size=40, + use_bias=True, + rnn_type="LSTM", + directions=directions, + ) + # Long sequence. + verify_rnn( + seq_length=8, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=True, + rnn_type="LSTM", + directions=directions, + ) + # Large hidden. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=128, + use_bias=True, + rnn_type="LSTM", + directions=directions, + ) + # Large input. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=64, + hidden_size=32, + use_bias=True, + rnn_type="LSTM", + directions=directions, + ) - # Testing for bidirectionality - verify_rnn( - seq_length=3, - batch_size=3, - input_size=16, - hidden_size=40, - use_bias=True, - rnn_type="LSTM", - directions=2, - ) + # Different activation testing. + # Default value hardsigmoid. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "Tanh", "Tanh"], + rnn_type="LSTM", + directions=directions, + ) + # Multiple parameterized activations. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "LeakyRelu", "Tanh"], + alphas=[2.0, 0.5], + betas=[0.3], + rnn_type="LSTM", + directions=directions, + ) + # All parameterized with new Affine activation. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "LeakyRelu", "Affine"], + alphas=[2.0, 0.5, 0.8], + betas=[0.3, 0.1], + rnn_type="LSTM", + directions=directions, + ) - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "LeakyRelu", "Affine"], - alphas=[2.0, 0.5, 0.8], - betas=[0.3, 0.1], - rnn_type="LSTM", - directions=2, - ) + # Testing with initial state and peepholes + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=True, + use_initial_state=True, + rnn_type="LSTM", + directions=directions, + ) - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=True, - use_initial_state=True, - use_peep=True, - rnn_type="LSTM", - directions=2, - ) + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=True, + use_initial_state=True, + use_peep=True, + rnn_type="LSTM", + directions=directions, + ) @tvm.testing.uses_gpu From c0d1b5fb6f3c153b08262fea11165668937bf3f2 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 15:34:01 -0700 Subject: [PATCH 17/27] double up activations in bidirect --- tests/python/frontend/onnx/test_forward.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 5a6a7f8bb864..68c3f1fc1232 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3371,7 +3371,7 @@ def test_lstm(): input_size=16, hidden_size=32, use_bias=False, - activations=["HardSigmoid", "Tanh", "Tanh"], + activations=["HardSigmoid", "Tanh", "Tanh"] * directions, rnn_type="LSTM", directions=directions, ) @@ -3382,7 +3382,7 @@ def test_lstm(): input_size=16, hidden_size=32, use_bias=False, - activations=["HardSigmoid", "LeakyRelu", "Tanh"], + activations=["HardSigmoid", "LeakyRelu", "Tanh"] * directions, alphas=[2.0, 0.5], betas=[0.3], rnn_type="LSTM", @@ -3395,7 +3395,7 @@ def test_lstm(): input_size=16, hidden_size=32, use_bias=False, - activations=["HardSigmoid", "LeakyRelu", "Affine"], + activations=["HardSigmoid", "LeakyRelu", "Affine"] * directions, alphas=[2.0, 0.5, 0.8], betas=[0.3, 0.1], rnn_type="LSTM", From 385970040a7d0d94b0d674421306ed49b409313c Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 15:36:49 -0700 Subject: [PATCH 18/27] slow attribute forgetting --- python/tvm/relay/frontend/onnx.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 0308b91abe20..64f39034f087 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -34,18 +34,9 @@ from .. import qnn as _qnn from .. import ty as _ty from .. import vision as _vision -from .common import ( - AttrCvt, - Renamer, - fold_constant, - get_name, - get_relay_op, - infer_channels, - infer_shape, - infer_type, - infer_value, - new_var, -) +from .common import (AttrCvt, Renamer, fold_constant, get_name, get_relay_op, + infer_channels, infer_shape, infer_type, infer_value, + new_var) __all__ = ["from_onnx"] @@ -2149,8 +2140,8 @@ def _impl_v7(cls, inputs, attr, params): if "activations" in attr: activations = attr["activations"] - if len(activations) != 3: - raise NotImplementedError("LSTM assumes 3 activation functions are provided") + if len(activations) != 3 * num_directions: + raise NotImplementedError(f"LSTM assumes 3 * num_directions activation functions are provided") alpha_loc = 0 alphas = attr.get("activation_alpha", []) if isinstance(alphas, float): From 5f3b4be2c65f2fae9be1376f189584e727c3b5dd Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 15:45:06 -0700 Subject: [PATCH 19/27] lstm interface is v. confus --- python/tvm/relay/frontend/onnx.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 64f39034f087..bfc1f8fbbdba 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2141,7 +2141,9 @@ def _impl_v7(cls, inputs, attr, params): if "activations" in attr: activations = attr["activations"] if len(activations) != 3 * num_directions: - raise NotImplementedError(f"LSTM assumes 3 * num_directions activation functions are provided") + raise NotImplementedError( + f"LSTM assumes 3 * num_directions activation functions are provided" + ) alpha_loc = 0 alphas = attr.get("activation_alpha", []) if isinstance(alphas, float): @@ -2151,7 +2153,7 @@ def _impl_v7(cls, inputs, attr, params): if isinstance(betas, float): betas = [betas] acts = [] - for i in range(3): + for i in range(3 * num_directions): alpha = None beta = None activation = activations[i] @@ -2162,11 +2164,8 @@ def _impl_v7(cls, inputs, attr, params): beta = betas[beta_loc] beta_loc += 1 acts.append(cls._activation_helper(activation, alpha, beta)) - f_act, g_act, h_act = acts else: - f_act = _op.sigmoid - g_act = _op.tanh - h_act = _op.tanh + acts = [_op.sigmoid, _op.tanh, _op.tanh] * num_directions X_steps = _op.split(X, indices_or_sections=X_shape[0], axis=0) result_output = [] @@ -2190,6 +2189,8 @@ def _impl_v7(cls, inputs, attr, params): p_i = _op.squeeze(p_is[i], axis=[0]) p_f = _op.squeeze(p_fs[i], axis=[0]) p_o = _op.squeeze(p_os[i], axis=[0]) + + f_act, g_act, h_act = acts[i * 3 : (i + 1) * 3] output, H, C = LSTM.generate_lstm( X_steps=X_steps, H_t=H_t, From 3b2e4865c20483d6c2ea965d9b18c0d3424b6f74 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 16:17:28 -0700 Subject: [PATCH 20/27] test forward complete --- tests/python/frontend/onnx/test_forward.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 68c3f1fc1232..b99358b3cde7 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3383,8 +3383,8 @@ def test_lstm(): hidden_size=32, use_bias=False, activations=["HardSigmoid", "LeakyRelu", "Tanh"] * directions, - alphas=[2.0, 0.5], - betas=[0.3], + alphas=[2.0, 0.5, 0] * directions, + betas=[0.3, 0, 0] * directions, rnn_type="LSTM", directions=directions, ) @@ -3396,8 +3396,8 @@ def test_lstm(): hidden_size=32, use_bias=False, activations=["HardSigmoid", "LeakyRelu", "Affine"] * directions, - alphas=[2.0, 0.5, 0.8], - betas=[0.3, 0.1], + alphas=[2.0, 0.5, 0.8] * directions, + betas=[0.3, 0.1, 0] * directions, rnn_type="LSTM", directions=directions, ) From fba5ba9cb156ef44d6fbbf3e4907bd7a385dfbd5 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 16:58:48 -0700 Subject: [PATCH 21/27] add GRU outline --- python/tvm/relay/frontend/onnx.py | 154 ++++++++++++++++++------------ 1 file changed, 95 insertions(+), 59 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index bfc1f8fbbdba..a3b7813e8e15 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2221,47 +2221,87 @@ def _impl_v7(cls, inputs, attr, params): class GRU(RNN): """Operator convert for GRU""" + @classmethod + def generate_gru( + cls, X_steps, H_t, W, R, B, linear_before_reset, f_act, g_act, W_dtype, backwards=False + ): + h_list = [] + seq_length = len(X_steps) + for i in range(seq_length): + step = X_steps[i] if not backwards else X_steps[seq_length - (i + 1)] + step = _op.squeeze(step, axis=[0]) + current = _op.nn.dense(step, W) + cz, cr, ch = _op.split(current, 3, axis=1) + rz, rr, rh = _op.split(R, 3, axis=0) + z = cz + _op.nn.dense(H_t, rz) + r = cr + _op.nn.dense(H_t, rr) + if B is not None: + WB, RB = _op.split(B, 2) + wbz, wbr, wbh = _op.split(WB, 3, axis=-1) + rbz, rbr, rbh = _op.split(RB, 3, axis=-1) + z += wbz + rbz + r += wbr + rbr + if linear_before_reset: + h = ch + (r * (_op.nn.dense(H_t, rh) + rbh)) + wbh + else: + h = ch + _op.nn.dense((r * H_t), rh) + wbh + rbh + else: + if linear_before_reset: + h = ch + (r * (_op.nn.dense(H_t, rh))) + else: + h = ch + _op.nn.dense((r * H_t), rh) + + z = f_act(z) + r = f_act(r) + h = g_act(h) + + H_t = ((_expr.const(1, dtype=W_dtype) - z) * h) + (z * H_t) + h_list.append(_op.expand_dims(H_t, axis=0)) + + # Concatenate outputs and add back in direction axis. + concatenated = _op.concatenate(h_list, 0) + output = _op.expand_dims(concatenated, axis=1) + H_t = _op.expand_dims(H_t, axis=0) + + return output, H_t + @classmethod def _impl_v7(cls, inputs, attr, params): # Unpack inputs, note that if optional and not provided then value will be None. X = inputs[0] - W = inputs[1] - R = inputs[2] - B = inputs[3] + Wp = inputs[1] + Rp = inputs[2] + Bp = inputs[3] # Sequence length currently unused as it can be inferred from shapes. # sequence_lens = inputs['sequence_lens'] - h_0 = inputs[5] + Hp_0 = inputs[5] linear_before_reset = attr.get("linear_before_reset", 0) - num_directions = infer_shape(W)[0] - W_dtype = infer_type(W).checked_type.dtype + num_directions = infer_shape(Wp)[0] + W_dtype = infer_type(Wp).checked_type.dtype - if num_directions != 1: - raise NotImplementedError("Bidirectional GRUs not yet supported.") - # Remove num_directions axis from weights. - W = _op.squeeze(W, axis=[0]) - R = _op.squeeze(R, axis=[0]) - if B is not None: - B = _op.squeeze(B, axis=[0]) + if num_directions not in [1, 2]: + raise NotImplementedError( + f"Directions for GRUs should be either 1 or 2 got {num_directions}" + ) X_shape = infer_shape(X) - hidden_size = infer_shape(R)[-1] + hidden_size = infer_shape(Rp)[-1] batch_size = X_shape[1] # Initialize state if not provided. # Otherwise remove bidirectional axis. - if h_0 is None: - h_0 = _op.zeros((batch_size, hidden_size), W_dtype) - else: - h_0 = _op.squeeze(h_0, axis=[0]) - - H_t = h_0 - h_list = [] + if Hp_0 is None: + Hp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) + if Bp is not None: + Bp = _op.zeros((num_directions, hidden_size * 6), W_dtype) if "activations" in attr: activations = attr["activations"] - if len(activations) != 2: - raise NotImplementedError("GRU assumes 2 activation functions are provided") + if len(activations) != 2 * num_directions: + raise NotImplementedError( + "GRU assumes 2 * num_directions activation functions are provided" + ) alpha_loc = 0 alphas = attr.get("activation_alpha", []) if isinstance(alphas, float): @@ -2282,47 +2322,43 @@ def _impl_v7(cls, inputs, attr, params): beta = betas[beta_loc] beta_loc += 1 acts.append(cls._activation_helper(activation, alpha, beta)) - f_act, g_act = acts else: - f_act = _op.sigmoid - g_act = _op.tanh + acts = [_op.sigmoid, _op.tanh] * 2 - X_steps = _op.split(X, indices_or_sections=X_shape[0], axis=0) - for step in X_steps: - step = _op.squeeze(step, axis=[0]) - current = _op.nn.dense(step, W) - cz, cr, ch = _op.split(current, 3, axis=1) - rz, rr, rh = _op.split(R, 3, axis=0) - z = cz + _op.nn.dense(H_t, rz) - r = cr + _op.nn.dense(H_t, rr) - if B is not None: - WB, RB = _op.split(B, 2) - wbz, wbr, wbh = _op.split(WB, 3, axis=-1) - rbz, rbr, rbh = _op.split(RB, 3, axis=-1) - z += wbz + rbz - r += wbr + rbr - if linear_before_reset: - h = ch + (r * (_op.nn.dense(H_t, rh) + rbh)) + wbh - else: - h = ch + _op.nn.dense((r * H_t), rh) + wbh + rbh - else: - if linear_before_reset: - h = ch + (r * (_op.nn.dense(H_t, rh))) - else: - h = ch + _op.nn.dense((r * H_t), rh) + result_output = [] + result_H = [] - z = f_act(z) - r = f_act(r) - h = g_act(h) + H_ts = _op.split(Hp_0, num_directions) + Ws = _op.split(Wp, num_directions) + Rs = _op.split(Rp, num_directions) + Bs = _op.split(Bp, num_directions) - H_t = ((_expr.const(1, dtype=W_dtype) - z) * h) + (z * H_t) - h_list.append(_op.expand_dims(H_t, axis=0)) - # Concatenate outputs and add back in direction axis. - concatenated = _op.concatenate(h_list, 0) - output = _op.expand_dims(concatenated, axis=1) - H_t = _op.expand_dims(H_t, axis=0) + for i in range(num_directions): + H_t = _op.squeeze(H_ts[i], axis=[0]) + W = _op.squeeze(Ws[i], axis=[0]) + R = _op.squeeze(Rs[i], axis=[0]) + B = _op.squeeze(Bs[i], axis=[0]) + f_act, g_act = acts[i * 2 : (i + 1) * 2] + output, H = GRU.generate_gru( + X_steps=X, + H_t=H_t, + W=W, + R=R, + B=B, + linear_before_reset=linear_before_reset, + f_act=f_act, + g_act=g_act, + W_dtype=W_dtype, + backwards=i == 1, + ) + + result_output.append(output) + result_H.append(H) + + output = _op.concatenate(result_output, axis=1) + H = _op.concatenate(result_H, axis=0) - return _expr.TupleWrapper(_expr.Tuple((output, H_t)), 2) + return _expr.TupleWrapper(_expr.Tuple((output, H)), 2) class Resize(OnnxOpConverter): From 4adca4b710c17ebc8e9dcb14d04c8eb97e884313 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 17:01:06 -0700 Subject: [PATCH 22/27] revisiion2 --- python/tvm/relay/frontend/onnx.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index a3b7813e8e15..a3081b811207 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -34,9 +34,18 @@ from .. import qnn as _qnn from .. import ty as _ty from .. import vision as _vision -from .common import (AttrCvt, Renamer, fold_constant, get_name, get_relay_op, - infer_channels, infer_shape, infer_type, infer_value, - new_var) +from .common import ( + AttrCvt, + Renamer, + fold_constant, + get_name, + get_relay_op, + infer_channels, + infer_shape, + infer_type, + infer_value, + new_var, +) __all__ = ["from_onnx"] @@ -2258,6 +2267,10 @@ def generate_gru( H_t = ((_expr.const(1, dtype=W_dtype) - z) * h) + (z * H_t) h_list.append(_op.expand_dims(H_t, axis=0)) + if backwards: + # Canonical view is hidden states from the first token not last + h_list = h_list[::-1] + # Concatenate outputs and add back in direction axis. concatenated = _op.concatenate(h_list, 0) output = _op.expand_dims(concatenated, axis=1) @@ -2328,6 +2341,7 @@ def _impl_v7(cls, inputs, attr, params): result_output = [] result_H = [] + X_steps = _op.split(X, indices_or_sections=X_shape[0], axis=0) H_ts = _op.split(Hp_0, num_directions) Ws = _op.split(Wp, num_directions) Rs = _op.split(Rp, num_directions) @@ -2340,7 +2354,7 @@ def _impl_v7(cls, inputs, attr, params): B = _op.squeeze(Bs[i], axis=[0]) f_act, g_act = acts[i * 2 : (i + 1) * 2] output, H = GRU.generate_gru( - X_steps=X, + X_steps=X_steps, H_t=H_t, W=W, R=R, From bf02a274a258cb7250373cdf2359399f813df7d7 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 17:03:44 -0700 Subject: [PATCH 23/27] why was tehre a not --- python/tvm/relay/frontend/onnx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index a3081b811207..631c46b66142 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2306,7 +2306,7 @@ def _impl_v7(cls, inputs, attr, params): # Otherwise remove bidirectional axis. if Hp_0 is None: Hp_0 = _op.zeros((num_directions, batch_size, hidden_size), W_dtype) - if Bp is not None: + if Bp is None: Bp = _op.zeros((num_directions, hidden_size * 6), W_dtype) if "activations" in attr: From 79544801ac4a63c26a2b902d87abdf3767afcaa6 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 17:06:09 -0700 Subject: [PATCH 24/27] gru tests --- tests/python/frontend/onnx/test_forward.py | 190 ++++++++++++--------- 1 file changed, 113 insertions(+), 77 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index b99358b3cde7..4f6f95615340 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3429,83 +3429,119 @@ def test_lstm(): @tvm.testing.uses_gpu def test_gru(): - # No bias. - verify_rnn( - seq_length=2, batch_size=1, input_size=16, hidden_size=32, use_bias=False, rnn_type="GRU" - ) - # large batch. - verify_rnn( - seq_length=4, - batch_size=8, - input_size=16, - hidden_size=32, - use_bias=True, - rnn_type="GRU", - linear_before_reset=True, - ) - # Non power of two. - verify_rnn( - seq_length=3, batch_size=3, input_size=16, hidden_size=40, use_bias=True, rnn_type="GRU" - ) - # Long sequence. - verify_rnn( - seq_length=8, batch_size=1, input_size=16, hidden_size=32, use_bias=True, rnn_type="GRU" - ) - # Large hidden. - verify_rnn( - seq_length=2, batch_size=1, input_size=16, hidden_size=128, use_bias=True, rnn_type="GRU" - ) - # Large input. - verify_rnn( - seq_length=2, batch_size=1, input_size=64, hidden_size=32, use_bias=True, rnn_type="GRU" - ) - - # Different activation testing. - # Default value hardsigmoid. - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "Softsign"], - rnn_type="GRU", - ) - # Multiple parameterized activations. - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "LeakyRelu"], - alphas=[2.0, 0.5], - betas=[0.3], - rnn_type="GRU", - ) - # All parameterized with new Affine activation. - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=False, - activations=["HardSigmoid", "Affine"], - alphas=[2.0, 0.8], - betas=[0.3, 0.1], - rnn_type="GRU", - ) - - # Testing with initial state - verify_rnn( - seq_length=2, - batch_size=1, - input_size=16, - hidden_size=32, - use_bias=True, - use_initial_state=True, - rnn_type="GRU", - ) + for directions in [1, 2]: + # No bias. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + rnn_type="GRU", + directions=directions, + ) + # large batch. + verify_rnn( + seq_length=4, + batch_size=8, + input_size=16, + hidden_size=32, + use_bias=True, + rnn_type="GRU", + linear_before_reset=True, + directions=directions, + ) + # Non power of two. + verify_rnn( + seq_length=3, + batch_size=3, + input_size=16, + hidden_size=40, + use_bias=True, + rnn_type="GRU", + directions=directions, + ) + # Long sequence. + verify_rnn( + seq_length=8, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=True, + rnn_type="GRU", + directions=directions, + ) + # Large hidden. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=128, + use_bias=True, + rnn_type="GRU", + directions=directions, + ) + # Large input. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=64, + hidden_size=32, + use_bias=True, + rnn_type="GRU", + directions=directions, + ) + + # Different activation testing. + # Default value hardsigmoid. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "Softsign"] * directions, + rnn_type="GRU", + directions=directions, + ) + # Multiple parameterized activations. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "LeakyRelu"] * directions, + alphas=[2.0, 0.5] * directions, + betas=[0.3, 0] * directions, + rnn_type="GRU", + directions=directions, + ) + # All parameterized with new Affine activation. + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=False, + activations=["HardSigmoid", "Affine"] * directions, + alphas=[2.0, 0.8] * directions, + betas=[0.3, 0.1] * directions, + rnn_type="GRU", + directions=directions, + ) + + # Testing with initial state + verify_rnn( + seq_length=2, + batch_size=1, + input_size=16, + hidden_size=32, + use_bias=True, + use_initial_state=True, + rnn_type="GRU", + directions=directions, + ) @tvm.testing.uses_gpu From d29011ad1e5321d1e0d6533febac4723c4315882 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 17:08:44 -0700 Subject: [PATCH 25/27] missing bounds, copy pasta! --- python/tvm/relay/frontend/onnx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 631c46b66142..692c236a8340 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2324,7 +2324,7 @@ def _impl_v7(cls, inputs, attr, params): if isinstance(betas, float): betas = [betas] acts = [] - for i in range(2): + for i in range(2 * num_directions): alpha = None beta = None activation = activations[i] From 855b54d436283bd1182a0f304b05be436b0b1273 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Thu, 24 Jun 2021 17:15:00 -0700 Subject: [PATCH 26/27] add comment --- python/tvm/relay/frontend/onnx.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 692c236a8340..b38ad332af82 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -2234,6 +2234,10 @@ class GRU(RNN): def generate_gru( cls, X_steps, H_t, W, R, B, linear_before_reset, f_act, g_act, W_dtype, backwards=False ): + """Create an unrolled gru loop. + + See https://github.com/onnx/onnx/blob/master/docs/Operators.md for math. + """ h_list = [] seq_length = len(X_steps) for i in range(seq_length): From 804bc7bd3012847a8da3ed1e9719cbb31a5a2441 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Fri, 25 Jun 2021 17:36:19 -0700 Subject: [PATCH 27/27] ensure all args fp --- tests/python/frontend/onnx/test_forward.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 4f6f95615340..db71855fd80f 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3383,8 +3383,8 @@ def test_lstm(): hidden_size=32, use_bias=False, activations=["HardSigmoid", "LeakyRelu", "Tanh"] * directions, - alphas=[2.0, 0.5, 0] * directions, - betas=[0.3, 0, 0] * directions, + alphas=[2.0, 0.5, 0.0] * directions, + betas=[0.3, 0.0, 0.0] * directions, rnn_type="LSTM", directions=directions, ) @@ -3397,7 +3397,7 @@ def test_lstm(): use_bias=False, activations=["HardSigmoid", "LeakyRelu", "Affine"] * directions, alphas=[2.0, 0.5, 0.8] * directions, - betas=[0.3, 0.1, 0] * directions, + betas=[0.3, 0.1, 0.0] * directions, rnn_type="LSTM", directions=directions, ) @@ -3513,7 +3513,7 @@ def test_gru(): use_bias=False, activations=["HardSigmoid", "LeakyRelu"] * directions, alphas=[2.0, 0.5] * directions, - betas=[0.3, 0] * directions, + betas=[0.3, 0.0] * directions, rnn_type="GRU", directions=directions, )