From 7d97224e745ab88dc7be5fedac4ba1ee6f39984f Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Fri, 26 Oct 2018 18:57:19 -0700 Subject: [PATCH 01/14] Port LSTM cell definition to Relay for testing --- python/tvm/relay/testing/layers.py | 4 +-- python/tvm/relay/testing/lstm.py | 47 ++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 python/tvm/relay/testing/lstm.py diff --git a/python/tvm/relay/testing/layers.py b/python/tvm/relay/testing/layers.py index 1b279d9e72af..9d4d3b3b4e13 100644 --- a/python/tvm/relay/testing/layers.py +++ b/python/tvm/relay/testing/layers.py @@ -105,7 +105,7 @@ def conv2d_transpose(data, weight=None, **kwargs): weight = relay.var(name + "_weight") return relay.nn.conv2d_transpose(data, weight, **kwargs) -def dense_add_bias(data, weight=None, bias=None, **kwargs): +def dense_add_bias(data, weight=None, bias=None, units=None, **kwargs): """Wrapper of dense which automatically creates weights if not given. Parameters @@ -133,6 +133,6 @@ def dense_add_bias(data, weight=None, bias=None, **kwargs): weight = relay.var(name + "_weight") if not bias: bias = relay.var(name + "_bias") - data = relay.nn.dense(data, weight, **kwargs) + data = relay.nn.dense(data, weight, units, **kwargs) data = relay.nn.bias_add(data, bias) return data diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py new file mode 100644 index 000000000000..d74a5a04e4a8 --- /dev/null +++ b/python/tvm/relay/testing/lstm.py @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Implementation of a Long Short-Term Memory (LSTM) cell. + +Adapted from: +https://gist.github.com/merrymercy/5eb24e3b019f84200645bd001e9caae9 +""" + +from tvm import realy +from . import layers +from .init import create_workload + +def lstm_cell(inputs, states, i2h_weight, h2h_weight, + i2h_bias, h2h_bias, num_hidden): + i2h = layers.dense_add_bias(data=inputs, weight=i2h_weight, + bias=i2h_bias, units=num_hidden * 4) + h2h = layers.dense_add_bias(data=inputs, weight=h2h_weight, + bias=h2h_bias, units=num_hidden * 4) + + gates = relay.add(i2h, h2h) + slice_gates = relay.split(gates, indices_or_sections=4) + + in_gate = relay.sigmoid(slice_gates[0]) + forget_gate = relay.sigmoid(slice_gates[1]) + in_transform = relay.tanh(slice_gates[2]) + out_gate = relay.sigmoid(slice_gates[3]) + next_c = relay.add(relay.mul(forget_gate, states[1]), + relay.mul(in_gate, in_transform)) + next_h = relay.mul(out_gate, relay.tanh(next_c)) + + return relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) From 3ca992ab1a64bde2df1e4b32d4913e9bb975a4c3 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Fri, 26 Oct 2018 19:05:38 -0700 Subject: [PATCH 02/14] Correct typo --- python/tvm/relay/testing/lstm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index d74a5a04e4a8..fed53091a029 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -22,7 +22,7 @@ https://gist.github.com/merrymercy/5eb24e3b019f84200645bd001e9caae9 """ -from tvm import realy +from tvm import relay from . import layers from .init import create_workload From 5dfe1147721780562eb08b05731178245c34fdbc Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Fri, 26 Oct 2018 19:15:43 -0700 Subject: [PATCH 03/14] Add docstring of questionable accuracy to LSTM --- python/tvm/relay/testing/lstm.py | 33 +++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index fed53091a029..cd912b07f403 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -28,9 +28,40 @@ def lstm_cell(inputs, states, i2h_weight, h2h_weight, i2h_bias, h2h_bias, num_hidden): + """Long-Short Term Memory (LSTM) network cell. + + Parameters + ---------- + inputs : relay.Expr + Input data to LSTM cell. + + states : Tuple[relay.Expr, relay.Expr] + State for LSTM cell. Should be of shape (batch size, num_hidden) + + i2h_weight : relay.Expr + i2h weight to LSTM cell. + + h2h_weight : relay.Expr + h2h weight to LSTM cell. + + i2h_bias : relay.Expr + i2h bias to LSTM cell. + + h2h_bias : relay.Expr + h2h bias to LSTM cell. + + num_hidden : int + Number of units in output symbol. + + Returns + ------- + result : Tuple[relay.Expr, Tuple[relay.Expr, relay.Expr]] + The result. + """ + i2h = layers.dense_add_bias(data=inputs, weight=i2h_weight, bias=i2h_bias, units=num_hidden * 4) - h2h = layers.dense_add_bias(data=inputs, weight=h2h_weight, + h2h = layers.dense_add_bias(data=states[0], weight=h2h_weight, bias=h2h_bias, units=num_hidden * 4) gates = relay.add(i2h, h2h) From e206c139e31a7c3a0ba8beea1147c5765c538c39 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Tue, 30 Oct 2018 15:01:51 -0700 Subject: [PATCH 04/14] Refactor LSTM so it returns a Relay function --- python/tvm/relay/testing/lstm.py | 53 +++++++++++++++----------------- 1 file changed, 24 insertions(+), 29 deletions(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index cd912b07f403..404649ed43a7 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -24,45 +24,37 @@ from tvm import relay from . import layers -from .init import create_workload +#from .init import create_workload -def lstm_cell(inputs, states, i2h_weight, h2h_weight, - i2h_bias, h2h_bias, num_hidden): +def lstm_cell(num_hidden, batch_size=1): """Long-Short Term Memory (LSTM) network cell. Parameters ---------- - inputs : relay.Expr - Input data to LSTM cell. - - states : Tuple[relay.Expr, relay.Expr] - State for LSTM cell. Should be of shape (batch size, num_hidden) - - i2h_weight : relay.Expr - i2h weight to LSTM cell. - - h2h_weight : relay.Expr - h2h weight to LSTM cell. - - i2h_bias : relay.Expr - i2h bias to LSTM cell. - - h2h_bias : relay.Expr - h2h bias to LSTM cell. - num_hidden : int Number of units in output symbol. + batch_size : int + Batch size (length of states). + Returns ------- - result : Tuple[relay.Expr, Tuple[relay.Expr, relay.Expr]] - The result. + result : relay.Expr + A Relay function that evaluates an LSTM cell. + The function takes in a tensor of input data, a tuple of two + states, and weights and biases for dense operations on the + inputs and on the state. It returns a tuple with two members, + an output tensor and a tuple of two new states. """ + inputs = relay.var("inputs") + states = relay.var("states", + relay.TupleType([ + relay.TensorType((batch_size, num_hidden)), + relay.TensorType((batch_size, num_hidden))])) - i2h = layers.dense_add_bias(data=inputs, weight=i2h_weight, - bias=i2h_bias, units=num_hidden * 4) - h2h = layers.dense_add_bias(data=states[0], weight=h2h_weight, - bias=h2h_bias, units=num_hidden * 4) + i2h = layers.dense_add_bias(data=inputs, units=num_hidden * 4) + h2h = layers.dense_add_bias(data=relay.TupleGetItem(states, 0), + units=num_hidden * 4) gates = relay.add(i2h, h2h) slice_gates = relay.split(gates, indices_or_sections=4) @@ -71,8 +63,11 @@ def lstm_cell(inputs, states, i2h_weight, h2h_weight, forget_gate = relay.sigmoid(slice_gates[1]) in_transform = relay.tanh(slice_gates[2]) out_gate = relay.sigmoid(slice_gates[3]) - next_c = relay.add(relay.mul(forget_gate, states[1]), + next_c = relay.add(relay.mul(forget_gate, + relay.TupleGetItem(states, 1)), relay.mul(in_gate, in_transform)) next_h = relay.mul(out_gate, relay.tanh(next_c)) + ret = relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) - return relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) + args = relay.ir_pass.free_vars(ret) + return relay.Function(args, ret) From 306cc01f0dd75f6b64ec04e13b87fdaaaf540835 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Tue, 30 Oct 2018 16:59:31 -0700 Subject: [PATCH 05/14] Build up LSTM cells into an RNN and add tests --- python/tvm/relay/testing/__init__.py | 1 + python/tvm/relay/testing/lstm.py | 119 +++++++++++++++++++-- tests/python/relay/test_ir_text_printer.py | 8 ++ 3 files changed, 117 insertions(+), 11 deletions(-) diff --git a/python/tvm/relay/testing/__init__.py b/python/tvm/relay/testing/__init__.py index 913f97ecd4a1..43160d64549c 100644 --- a/python/tvm/relay/testing/__init__.py +++ b/python/tvm/relay/testing/__init__.py @@ -6,4 +6,5 @@ from . import dqn from . import dcgan from . import mobilenet +from . import lstm from .config import ctx_list diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index 404649ed43a7..63de5d5d3d4c 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -24,9 +24,9 @@ from tvm import relay from . import layers -#from .init import create_workload +from .init import create_workload -def lstm_cell(num_hidden, batch_size=1): +def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): """Long-Short Term Memory (LSTM) network cell. Parameters @@ -39,7 +39,7 @@ def lstm_cell(num_hidden, batch_size=1): Returns ------- - result : relay.Expr + result : tvm.relay.Function A Relay function that evaluates an LSTM cell. The function takes in a tensor of input data, a tuple of two states, and weights and biases for dense operations on the @@ -49,12 +49,14 @@ def lstm_cell(num_hidden, batch_size=1): inputs = relay.var("inputs") states = relay.var("states", relay.TupleType([ - relay.TensorType((batch_size, num_hidden)), - relay.TensorType((batch_size, num_hidden))])) + relay.TensorType((batch_size, num_hidden), dtype), + relay.TensorType((batch_size, num_hidden), dtype)])) - i2h = layers.dense_add_bias(data=inputs, units=num_hidden * 4) + i2h = layers.dense_add_bias(data=inputs, units=num_hidden * 4, + name="%si2h" % name) h2h = layers.dense_add_bias(data=relay.TupleGetItem(states, 0), - units=num_hidden * 4) + units=num_hidden * 4, + name="%sh2h" % name) gates = relay.add(i2h, h2h) slice_gates = relay.split(gates, indices_or_sections=4) @@ -63,11 +65,106 @@ def lstm_cell(num_hidden, batch_size=1): forget_gate = relay.sigmoid(slice_gates[1]) in_transform = relay.tanh(slice_gates[2]) out_gate = relay.sigmoid(slice_gates[3]) - next_c = relay.add(relay.mul(forget_gate, - relay.TupleGetItem(states, 1)), - relay.mul(in_gate, in_transform)) - next_h = relay.mul(out_gate, relay.tanh(next_c)) + next_c = relay.add(relay.multiply(forget_gate, + relay.TupleGetItem(states, 1)), + relay.multiply(in_gate, in_transform)) + next_h = relay.multiply(out_gate, relay.tanh(next_c)) ret = relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) args = relay.ir_pass.free_vars(ret) return relay.Function(args, ret) + + +def rnn_builder(iterations, cell_fn, init_states, out, forward): + """Recursive builder of unrolled RNN: Returns let-chain of cell function calls. + + Parameters + ---------- + iterations : int + Number of iterations for the unrolled RNN. + + cell_fn : tvm.relay.Function + Relay function implementing the RNN cell computation. + + init_states : tvm.relay.Expr + A Relay tuple representing the state for the first iteration. + + out : tvm.relay.Var + A Relay variable into which the last iteration should assign + its output. + + forward : tvm.relay.Expr + A Relay AST that uses the result of the last iteration's cell. + + Returns + ------- + ret : tvm.relay.Expr + A let-chain of cell function calls of the number of iterations, + with the last iteration using the "forward" argument as the + output of the let chain. + """ + i = iterations + inputs = relay.Var("inputs_%s" % i) + i2h_weight = relay.Var("i2h_%s_weight" % i) + i2h_bias = relay.Var("i2h_%i_bias" % i) + h2h_weight = relay.Var("h2h_%s_weight" % i) + h2h_bias = relay.Var("h2h_%s_bias" % i) + + # base case: 0 is the first iteration, so use initial state + if i == 0: + return relay.Let(out, + relay.Call(cell_fn, + [inputs, init_states, i2h_weight, + i2h_bias, h2h_weight, h2h_bias]), + forward) + + # otherwise: create the chain backwards and insert in the last iteration + prev_out = relay.Var("out_%s" % (i - 1)) + call = relay.Let(out, + relay.Call(cell_fn, + [inputs, + relay.TupleGetItem(prev_out, 1), + i2h_bias, i2h_weight, + h2h_bias, h2h_weight]), + forward) + return rnn_builder(i - 1, cell_fn, init_states, prev_out, call) + + +def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): + '''Constructs an unrolled RNN with LSTM cells''' + states = relay.Tuple([relay.zeros((batch_size, num_hidden), dtype), + relay.zeros((batch_size, num_hidden), dtype)]) + + cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm") + + out = relay.Var("lstm_out") + get_value = relay.TupleGetItem(out, 0) + unrolled = rnn_builder(iterations - 1, cell_fn, states, + out, get_value) + + args = relay.ir_pass.free_vars(unrolled) + return relay.Function(args, unrolled) + + +def get_workload(iterations, num_hidden, batch_size=1, dtype="float32"): + """Get benchmark workload for an LSTM RNN. + + Parameters + ---------- + iterations : int + The number of iterations in the desired LSTM RNN. + num_hidden : int + The size of the hiddxen state + batch_size : int, optional (default 1) + The batch size used in the model + dtype : str, optional (default "float32") + The data type + Returns + ------- + net : nnvm.symbol + The computational graph + params : dict of str to NDArray + The parameters. + """ + net = get_net(batch_size, num_hidden, batch_size, dtype) + return create_workload(net) diff --git a/tests/python/relay/test_ir_text_printer.py b/tests/python/relay/test_ir_text_printer.py index dd790a6d7d87..fe30ffbcf974 100644 --- a/tests/python/relay/test_ir_text_printer.py +++ b/tests/python/relay/test_ir_text_printer.py @@ -96,10 +96,12 @@ def test_variable_name(): v1 = relay.var("1") assert "%v1" in v1.astext() + def test_mlp(): net, params = tvm.relay.testing.mlp.get_workload(batch_size=1) net.astext() + def test_resnet(): net, params = tvm.relay.testing.resnet.get_workload(batch_size=1) net.astext() @@ -116,6 +118,12 @@ def test_dcgan(): net, params = tvm.relay.testing.dcgan.get_workload(batch_size=1) net.astext() + +def test_lstm(): + net, params = tvm.relay.testing.lstm.get_workload(4, 4) + show(net.astext()) + + if __name__ == "__main__": do_print[0] = True test_resnet() From e4c626f3bb18723850b349ea080e843cc856c384 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Tue, 30 Oct 2018 19:18:29 -0700 Subject: [PATCH 06/14] Annotate all types in LSTM in hopes of getting it to typecheck (it didn't) --- python/tvm/relay/testing/lstm.py | 346 ++++++++++++++++--------------- 1 file changed, 176 insertions(+), 170 deletions(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index 63de5d5d3d4c..253ffc66adf5 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -1,170 +1,176 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -""" -Implementation of a Long Short-Term Memory (LSTM) cell. - -Adapted from: -https://gist.github.com/merrymercy/5eb24e3b019f84200645bd001e9caae9 -""" - -from tvm import relay -from . import layers -from .init import create_workload - -def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): - """Long-Short Term Memory (LSTM) network cell. - - Parameters - ---------- - num_hidden : int - Number of units in output symbol. - - batch_size : int - Batch size (length of states). - - Returns - ------- - result : tvm.relay.Function - A Relay function that evaluates an LSTM cell. - The function takes in a tensor of input data, a tuple of two - states, and weights and biases for dense operations on the - inputs and on the state. It returns a tuple with two members, - an output tensor and a tuple of two new states. - """ - inputs = relay.var("inputs") - states = relay.var("states", - relay.TupleType([ - relay.TensorType((batch_size, num_hidden), dtype), - relay.TensorType((batch_size, num_hidden), dtype)])) - - i2h = layers.dense_add_bias(data=inputs, units=num_hidden * 4, - name="%si2h" % name) - h2h = layers.dense_add_bias(data=relay.TupleGetItem(states, 0), - units=num_hidden * 4, - name="%sh2h" % name) - - gates = relay.add(i2h, h2h) - slice_gates = relay.split(gates, indices_or_sections=4) - - in_gate = relay.sigmoid(slice_gates[0]) - forget_gate = relay.sigmoid(slice_gates[1]) - in_transform = relay.tanh(slice_gates[2]) - out_gate = relay.sigmoid(slice_gates[3]) - next_c = relay.add(relay.multiply(forget_gate, - relay.TupleGetItem(states, 1)), - relay.multiply(in_gate, in_transform)) - next_h = relay.multiply(out_gate, relay.tanh(next_c)) - ret = relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) - - args = relay.ir_pass.free_vars(ret) - return relay.Function(args, ret) - - -def rnn_builder(iterations, cell_fn, init_states, out, forward): - """Recursive builder of unrolled RNN: Returns let-chain of cell function calls. - - Parameters - ---------- - iterations : int - Number of iterations for the unrolled RNN. - - cell_fn : tvm.relay.Function - Relay function implementing the RNN cell computation. - - init_states : tvm.relay.Expr - A Relay tuple representing the state for the first iteration. - - out : tvm.relay.Var - A Relay variable into which the last iteration should assign - its output. - - forward : tvm.relay.Expr - A Relay AST that uses the result of the last iteration's cell. - - Returns - ------- - ret : tvm.relay.Expr - A let-chain of cell function calls of the number of iterations, - with the last iteration using the "forward" argument as the - output of the let chain. - """ - i = iterations - inputs = relay.Var("inputs_%s" % i) - i2h_weight = relay.Var("i2h_%s_weight" % i) - i2h_bias = relay.Var("i2h_%i_bias" % i) - h2h_weight = relay.Var("h2h_%s_weight" % i) - h2h_bias = relay.Var("h2h_%s_bias" % i) - - # base case: 0 is the first iteration, so use initial state - if i == 0: - return relay.Let(out, - relay.Call(cell_fn, - [inputs, init_states, i2h_weight, - i2h_bias, h2h_weight, h2h_bias]), - forward) - - # otherwise: create the chain backwards and insert in the last iteration - prev_out = relay.Var("out_%s" % (i - 1)) - call = relay.Let(out, - relay.Call(cell_fn, - [inputs, - relay.TupleGetItem(prev_out, 1), - i2h_bias, i2h_weight, - h2h_bias, h2h_weight]), - forward) - return rnn_builder(i - 1, cell_fn, init_states, prev_out, call) - - -def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): - '''Constructs an unrolled RNN with LSTM cells''' - states = relay.Tuple([relay.zeros((batch_size, num_hidden), dtype), - relay.zeros((batch_size, num_hidden), dtype)]) - - cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm") - - out = relay.Var("lstm_out") - get_value = relay.TupleGetItem(out, 0) - unrolled = rnn_builder(iterations - 1, cell_fn, states, - out, get_value) - - args = relay.ir_pass.free_vars(unrolled) - return relay.Function(args, unrolled) - - -def get_workload(iterations, num_hidden, batch_size=1, dtype="float32"): - """Get benchmark workload for an LSTM RNN. - - Parameters - ---------- - iterations : int - The number of iterations in the desired LSTM RNN. - num_hidden : int - The size of the hiddxen state - batch_size : int, optional (default 1) - The batch size used in the model - dtype : str, optional (default "float32") - The data type - Returns - ------- - net : nnvm.symbol - The computational graph - params : dict of str to NDArray - The parameters. - """ - net = get_net(batch_size, num_hidden, batch_size, dtype) - return create_workload(net) +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Implementation of a Long Short-Term Memory (LSTM) cell. + +Adapted from: +https://gist.github.com/merrymercy/5eb24e3b019f84200645bd001e9caae9 +""" + +from tvm import relay +from . import layers +from .init import create_workload + +def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): + """Long-Short Term Memory (LSTM) network cell. + + Parameters + ---------- + num_hidden : int + Number of units in output symbol. + + batch_size : int + Batch size (length of states). + + Returns + ------- + result : tvm.relay.Function + A Relay function that evaluates an LSTM cell. + The function takes in a tensor of input data, a tuple of two + states, and weights and biases for dense operations on the + inputs and on the state. It returns a tuple with two members, + an output tensor and a tuple of two new states. + """ + input_type = relay.TensorType((batch_size, num_hidden), dtype) + weight_type = relay.TensorType((num_hidden, 4), dtype) + bias_type = relay.TensorType((4,), dtype) + + inputs = relay.Var("inputs", input_type) + states = relay.Var("states", + relay.TupleType([input_type, input_type])) + + i2h_weight = relay.Var("i2h_weight", weight_type) + i2h_bias = relay.Var("i2h_bias", bias_type) + + h2h_weight = relay.Var("h2h_weight", weight_type) + h2h_bias = relay.Var("h2h_bias", bias_type) + + i2h = layers.dense_add_bias(data=inputs, units=num_hidden * 4, + weight=i2h_weight, bias=i2h_bias, + name="%si2h" % name) + h2h = layers.dense_add_bias(data=relay.TupleGetItem(states, 0), + units=num_hidden * 4, + weight=h2h_weight, bias=h2h_bias, + name="%sh2h" % name) + + gates = relay.add(i2h, h2h) + slice_gates = relay.split(gates, indices_or_sections=4, axis=1) + + in_gate = relay.sigmoid(slice_gates[0]) + forget_gate = relay.sigmoid(slice_gates[1]) + in_transform = relay.tanh(slice_gates[2]) + out_gate = relay.sigmoid(slice_gates[3]) + next_c = relay.add(relay.multiply(forget_gate, + relay.TupleGetItem(states, 1)), + relay.multiply(in_gate, in_transform)) + next_h = relay.multiply(out_gate, relay.tanh(next_c)) + ret = relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) + + return relay.Function([inputs, states, i2h_weight, + i2h_bias, h2h_weight, h2h_bias], + ret, + relay.TupleType([ + input_type, + relay.TupleType([input_type, + input_type])])) + + +def rnn_builder(iterations, num_hidden, batch_size, dtype, out, forward): + """Recursive builder of unrolled RNN: Returns let-chain of cell function calls. + """ + i = iterations + + input_type = relay.TensorType((batch_size, num_hidden), dtype) + weight_type = relay.TensorType((num_hidden, 4), dtype) + bias_type = relay.TensorType((4,), dtype) + + inputs = relay.Var("inputs_%s" % i, input_type) + i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type) + i2h_bias = relay.Var("i2h_%i_bias" % i, bias_type) + h2h_weight = relay.Var("h2h_%s_weight" % i, weight_type) + h2h_bias = relay.Var("h2h_%s_bias" % i, bias_type) + + cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_%s" % i) + + # base case: 0 is the first iteration, so use initial state + if i == 0: + return relay.Let(out, + relay.Call(cell_fn, + [inputs, + relay.Tuple([ + relay.zeros((batch_size, num_hidden), dtype), + relay.zeros((batch_size, num_hidden), dtype) + ]), + i2h_weight, i2h_bias, + h2h_weight, h2h_bias]), + forward) + + # otherwise: create the chain backwards and insert in the last iteration + prev_out = relay.Var("out_%s" % (i - 1), + relay.TupleType([input_type, + relay.TupleType([input_type, + input_type])])) + call = relay.Let(out, + relay.Call(cell_fn, + [inputs, + relay.TupleGetItem(prev_out, 1), + i2h_weight, i2h_bias, + h2h_weight, h2h_bias]), + forward) + return rnn_builder(i - 1, num_hidden, batch_size, dtype, + prev_out, call) + + +def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): + '''Constructs an unrolled RNN with LSTM cells''' + input_type = relay.TensorType((batch_size, num_hidden), dtype) + out = relay.Var("lstm_out", + relay.TupleType([input_type, + relay.TupleType([input_type, + input_type])])) + get_value = relay.TupleGetItem(out, 0) + unrolled = rnn_builder(iterations - 1, + num_hidden, batch_size, dtype, + out, get_value) + + args = relay.ir_pass.free_vars(unrolled) + return relay.Function(args, unrolled, input_type) + + +def get_workload(iterations, num_hidden, batch_size=1, dtype="float32"): + """Get benchmark workload for an LSTM RNN. + + Parameters + ---------- + iterations : int + The number of iterations in the desired LSTM RNN. + num_hidden : int + The size of the hiddxen state + batch_size : int, optional (default 1) + The batch size used in the model + dtype : str, optional (default "float32") + The data type + Returns + ------- + net : nnvm.symbol + The computational graph + params : dict of str to NDArray + The parameters. + """ + net = get_net(iterations, num_hidden, batch_size, dtype) + return create_workload(net) From e70089ec519cea71047196d029ffd2ab69911c91 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Tue, 30 Oct 2018 19:35:30 -0700 Subject: [PATCH 07/14] Correct the weight and bias shapes --- python/tvm/relay/testing/lstm.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index 253ffc66adf5..8b4b352f9bd1 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -47,8 +47,8 @@ def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): an output tensor and a tuple of two new states. """ input_type = relay.TensorType((batch_size, num_hidden), dtype) - weight_type = relay.TensorType((num_hidden, 4), dtype) - bias_type = relay.TensorType((4,), dtype) + weight_type = relay.TensorType((num_hidden, 4*num_hidden), dtype) + bias_type = relay.TensorType((4*num_hidden,), dtype) inputs = relay.Var("inputs", input_type) states = relay.Var("states", @@ -96,8 +96,8 @@ def rnn_builder(iterations, num_hidden, batch_size, dtype, out, forward): i = iterations input_type = relay.TensorType((batch_size, num_hidden), dtype) - weight_type = relay.TensorType((num_hidden, 4), dtype) - bias_type = relay.TensorType((4,), dtype) + weight_type = relay.TensorType((num_hidden, 4*num_hidden), dtype) + bias_type = relay.TensorType((4*num_hidden,), dtype) inputs = relay.Var("inputs_%s" % i, input_type) i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type) From 058d4387e0017efb14a4f057b8d8911cb763e81c Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Wed, 31 Oct 2018 15:59:02 -0700 Subject: [PATCH 08/14] Use ScopeBuilder to annotate types at every step in LSTM --- python/tvm/relay/testing/lstm.py | 168 +++++++++++---------- tests/python/relay/test_ir_text_printer.py | 2 +- 2 files changed, 88 insertions(+), 82 deletions(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index 8b4b352f9bd1..6c951aebbeef 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -46,10 +46,18 @@ def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): inputs and on the state. It returns a tuple with two members, an output tensor and a tuple of two new states. """ + sb = relay.ScopeBuilder() + input_type = relay.TensorType((batch_size, num_hidden), dtype) weight_type = relay.TensorType((num_hidden, 4*num_hidden), dtype) bias_type = relay.TensorType((4*num_hidden,), dtype) + dense_type = relay.TensorType((batch_size, 4*num_hidden), dtype) + slice_type = relay.TupleType([input_type, input_type, + input_type, input_type]) + ret_type = relay.TupleType([input_type, + relay.TupleType([input_type, input_type])]) + inputs = relay.Var("inputs", input_type) states = relay.Var("states", relay.TupleType([input_type, input_type])) @@ -60,96 +68,94 @@ def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): h2h_weight = relay.Var("h2h_weight", weight_type) h2h_bias = relay.Var("h2h_bias", bias_type) - i2h = layers.dense_add_bias(data=inputs, units=num_hidden * 4, - weight=i2h_weight, bias=i2h_bias, - name="%si2h" % name) - h2h = layers.dense_add_bias(data=relay.TupleGetItem(states, 0), - units=num_hidden * 4, - weight=h2h_weight, bias=h2h_bias, - name="%sh2h" % name) - - gates = relay.add(i2h, h2h) - slice_gates = relay.split(gates, indices_or_sections=4, axis=1) - - in_gate = relay.sigmoid(slice_gates[0]) - forget_gate = relay.sigmoid(slice_gates[1]) - in_transform = relay.tanh(slice_gates[2]) - out_gate = relay.sigmoid(slice_gates[3]) - next_c = relay.add(relay.multiply(forget_gate, - relay.TupleGetItem(states, 1)), - relay.multiply(in_gate, in_transform)) - next_h = relay.multiply(out_gate, relay.tanh(next_c)) - ret = relay.Tuple([next_h, relay.Tuple([next_h, next_c])]) + i2h = sb.let(("i2h", dense_type), + layers.dense_add_bias( + data=inputs, + units=num_hidden * 4, + weight=i2h_weight, bias=i2h_bias, + name="%si2h" % name)) + h2h = sb.let(("h2h", dense_type), + layers.dense_add_bias( + data=relay.TupleGetItem(states, 0), + units=num_hidden * 4, + weight=h2h_weight, bias=h2h_bias, + name="%sh2h" % name)) + + gates = sb.let(("gates", dense_type), relay.add(i2h, h2h)) + slice_gates = sb.let(("slice_gates", slice_type), + relay.split(gates, + indices_or_sections=4, + axis=1).astuple()) + + in_gate = sb.let(("in_gate", input_type), + relay.sigmoid(relay.TupleGetItem(slice_gates, 0))) + forget_gate = sb.let(("forget_gate", input_type), + relay.sigmoid(relay.TupleGetItem(slice_gates, 1))) + in_transform = sb.let(("in_transform", input_type), + relay.tanh(relay.TupleGetItem(slice_gates, 2))) + out_gate = sb.let(("out_gate", input_type), + relay.sigmoid(relay.TupleGetItem(slice_gates, 3))) + + next_c = sb.let(("next_c", input_type), + relay.add(relay.multiply(forget_gate, + relay.TupleGetItem(states, 1)), + relay.multiply(in_gate, in_transform))) + next_h = sb.let(("next_h", input_type), + relay.multiply(out_gate, relay.tanh(next_c))) + ret = sb.let(("ret", ret_type), + relay.Tuple([next_h, relay.Tuple([next_h, next_c])])) + sb.ret(ret) + + body = sb.get() return relay.Function([inputs, states, i2h_weight, i2h_bias, h2h_weight, h2h_bias], - ret, - relay.TupleType([ - input_type, - relay.TupleType([input_type, - input_type])])) - + body, ret_type) -def rnn_builder(iterations, num_hidden, batch_size, dtype, out, forward): - """Recursive builder of unrolled RNN: Returns let-chain of cell function calls. - """ - i = iterations +def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): + '''Constructs an unrolled RNN with LSTM cells''' input_type = relay.TensorType((batch_size, num_hidden), dtype) weight_type = relay.TensorType((num_hidden, 4*num_hidden), dtype) bias_type = relay.TensorType((4*num_hidden,), dtype) - inputs = relay.Var("inputs_%s" % i, input_type) - i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type) - i2h_bias = relay.Var("i2h_%i_bias" % i, bias_type) - h2h_weight = relay.Var("h2h_%s_weight" % i, weight_type) - h2h_bias = relay.Var("h2h_%s_bias" % i, bias_type) - - cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_%s" % i) - - # base case: 0 is the first iteration, so use initial state - if i == 0: - return relay.Let(out, - relay.Call(cell_fn, - [inputs, - relay.Tuple([ - relay.zeros((batch_size, num_hidden), dtype), - relay.zeros((batch_size, num_hidden), dtype) - ]), - i2h_weight, i2h_bias, - h2h_weight, h2h_bias]), - forward) - - # otherwise: create the chain backwards and insert in the last iteration - prev_out = relay.Var("out_%s" % (i - 1), - relay.TupleType([input_type, - relay.TupleType([input_type, - input_type])])) - call = relay.Let(out, - relay.Call(cell_fn, - [inputs, - relay.TupleGetItem(prev_out, 1), - i2h_weight, i2h_bias, - h2h_weight, h2h_bias]), - forward) - return rnn_builder(i - 1, num_hidden, batch_size, dtype, - prev_out, call) - - -def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): - '''Constructs an unrolled RNN with LSTM cells''' - input_type = relay.TensorType((batch_size, num_hidden), dtype) - out = relay.Var("lstm_out", - relay.TupleType([input_type, - relay.TupleType([input_type, - input_type])])) - get_value = relay.TupleGetItem(out, 0) - unrolled = rnn_builder(iterations - 1, - num_hidden, batch_size, dtype, - out, get_value) - - args = relay.ir_pass.free_vars(unrolled) - return relay.Function(args, unrolled, input_type) + state_type = relay.TupleType([input_type, input_type]) + cell_type = relay.TupleType([input_type, state_type]) + + sb = relay.ScopeBuilder() + + zeros = sb.let(("zeros", input_type), + relay.zeros((batch_size, num_hidden), dtype)) + init_states = sb.let(("init_states", state_type), + relay.Tuple([zeros, zeros])) + + states = init_states + out = None + + for i in range(iterations): + inputs = relay.Var("inputs_%s" % i, input_type) + i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type) + i2h_bias = relay.Var("i2h_%i_bias" % i, bias_type) + h2h_weight = relay.Var("h2h_%s_weight" % i, weight_type) + h2h_bias = relay.Var("h2h_%s_bias" % i, bias_type) + + cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_%s" % i) + + call = sb.let(("call_%s" % i, cell_type), + relay.Call(cell_fn, + [inputs, states, i2h_weight, + i2h_bias, h2h_weight, h2h_bias])) + new_out = sb.let(("out_%s" % i, input_type), + relay.TupleGetItem(call, 0)) + new_states = sb.let(("states_%s" % i, state_type), + relay.TupleGetItem(call, 1)) + states = new_states + out = new_out + + sb.ret(out) + body = sb.get() + args = relay.ir_pass.free_vars(body) + return relay.Function(args, body, input_type) def get_workload(iterations, num_hidden, batch_size=1, dtype="float32"): diff --git a/tests/python/relay/test_ir_text_printer.py b/tests/python/relay/test_ir_text_printer.py index fe30ffbcf974..0ce36eafe612 100644 --- a/tests/python/relay/test_ir_text_printer.py +++ b/tests/python/relay/test_ir_text_printer.py @@ -121,7 +121,7 @@ def test_dcgan(): def test_lstm(): net, params = tvm.relay.testing.lstm.get_workload(4, 4) - show(net.astext()) + net.astext() if __name__ == "__main__": From 37cfed76b8e4b21c57c4739179c2008e89a9bcf2 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Wed, 31 Oct 2018 16:06:59 -0700 Subject: [PATCH 09/14] Rename ScopeBuilder to please pylint --- python/tvm/relay/testing/lstm.py | 108 +++++++++++++++---------------- 1 file changed, 54 insertions(+), 54 deletions(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index 6c951aebbeef..7bc777858008 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -68,45 +68,45 @@ def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): h2h_weight = relay.Var("h2h_weight", weight_type) h2h_bias = relay.Var("h2h_bias", bias_type) - i2h = sb.let(("i2h", dense_type), - layers.dense_add_bias( - data=inputs, - units=num_hidden * 4, - weight=i2h_weight, bias=i2h_bias, - name="%si2h" % name)) - h2h = sb.let(("h2h", dense_type), - layers.dense_add_bias( - data=relay.TupleGetItem(states, 0), - units=num_hidden * 4, - weight=h2h_weight, bias=h2h_bias, - name="%sh2h" % name)) - - gates = sb.let(("gates", dense_type), relay.add(i2h, h2h)) - slice_gates = sb.let(("slice_gates", slice_type), - relay.split(gates, - indices_or_sections=4, - axis=1).astuple()) - - in_gate = sb.let(("in_gate", input_type), - relay.sigmoid(relay.TupleGetItem(slice_gates, 0))) - forget_gate = sb.let(("forget_gate", input_type), - relay.sigmoid(relay.TupleGetItem(slice_gates, 1))) - in_transform = sb.let(("in_transform", input_type), - relay.tanh(relay.TupleGetItem(slice_gates, 2))) - out_gate = sb.let(("out_gate", input_type), - relay.sigmoid(relay.TupleGetItem(slice_gates, 3))) - - next_c = sb.let(("next_c", input_type), - relay.add(relay.multiply(forget_gate, - relay.TupleGetItem(states, 1)), - relay.multiply(in_gate, in_transform))) - next_h = sb.let(("next_h", input_type), - relay.multiply(out_gate, relay.tanh(next_c))) - ret = sb.let(("ret", ret_type), - relay.Tuple([next_h, relay.Tuple([next_h, next_c])])) - sb.ret(ret) - - body = sb.get() + i2h = builder.let(("i2h", dense_type), + layers.dense_add_bias( + data=inputs, + units=num_hidden * 4, + weight=i2h_weight, bias=i2h_bias, + name="%si2h" % name)) + h2h = builder.let(("h2h", dense_type), + layers.dense_add_bias( + data=relay.TupleGetItem(states, 0), + units=num_hidden * 4, + weight=h2h_weight, bias=h2h_bias, + name="%sh2h" % name)) + + gates = builder.let(("gates", dense_type), relay.add(i2h, h2h)) + slice_gates = builder.let(("slice_gates", slice_type), + relay.split(gates, + indices_or_sections=4, + axis=1).astuple()) + + in_gate = builder.let(("in_gate", input_type), + relay.sigmoid(relay.TupleGetItem(slice_gates, 0))) + forget_gate = builder.let(("forget_gate", input_type), + relay.sigmoid(relay.TupleGetItem(slice_gates, 1))) + in_transform = builder.let(("in_transform", input_type), + relay.tanh(relay.TupleGetItem(slice_gates, 2))) + out_gate = builder.let(("out_gate", input_type), + relay.sigmoid(relay.TupleGetItem(slice_gates, 3))) + + next_c = builder.let(("next_c", input_type), + relay.add(relay.multiply(forget_gate, + relay.TupleGetItem(states, 1)), + relay.multiply(in_gate, in_transform))) + next_h = builder.let(("next_h", input_type), + relay.multiply(out_gate, relay.tanh(next_c))) + ret = builder.let(("ret", ret_type), + relay.Tuple([next_h, relay.Tuple([next_h, next_c])])) + builder.ret(ret) + + body = builder.get() return relay.Function([inputs, states, i2h_weight, i2h_bias, h2h_weight, h2h_bias], @@ -122,12 +122,12 @@ def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): state_type = relay.TupleType([input_type, input_type]) cell_type = relay.TupleType([input_type, state_type]) - sb = relay.ScopeBuilder() + builder = relay.ScopeBuilder() - zeros = sb.let(("zeros", input_type), - relay.zeros((batch_size, num_hidden), dtype)) - init_states = sb.let(("init_states", state_type), - relay.Tuple([zeros, zeros])) + zeros = builder.let(("zeros", input_type), + relay.zeros((batch_size, num_hidden), dtype)) + init_states = builder.let(("init_states", state_type), + relay.Tuple([zeros, zeros])) states = init_states out = None @@ -141,19 +141,19 @@ def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_%s" % i) - call = sb.let(("call_%s" % i, cell_type), - relay.Call(cell_fn, - [inputs, states, i2h_weight, - i2h_bias, h2h_weight, h2h_bias])) - new_out = sb.let(("out_%s" % i, input_type), - relay.TupleGetItem(call, 0)) - new_states = sb.let(("states_%s" % i, state_type), - relay.TupleGetItem(call, 1)) + call = builder.let(("call_%s" % i, cell_type), + relay.Call(cell_fn, + [inputs, states, i2h_weight, + i2h_bias, h2h_weight, h2h_bias])) + new_out = builder.let(("out_%s" % i, input_type), + relay.TupleGetItem(call, 0)) + new_states = builder.let(("states_%s" % i, state_type), + relay.TupleGetItem(call, 1)) states = new_states out = new_out - sb.ret(out) - body = sb.get() + builder.ret(out) + body = builder.get() args = relay.ir_pass.free_vars(body) return relay.Function(args, body, input_type) From ae060b83c0c2e5780e47039c7200a94fe21a02fe Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Wed, 31 Oct 2018 16:09:46 -0700 Subject: [PATCH 10/14] Missed one --- python/tvm/relay/testing/lstm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index 7bc777858008..bc3991cb2434 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -46,7 +46,7 @@ def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""): inputs and on the state. It returns a tuple with two members, an output tensor and a tuple of two new states. """ - sb = relay.ScopeBuilder() + builder = relay.ScopeBuilder() input_type = relay.TensorType((batch_size, num_hidden), dtype) weight_type = relay.TensorType((num_hidden, 4*num_hidden), dtype) From c63b2862b108eccd273ed097198d044f200271f2 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Wed, 31 Oct 2018 16:37:22 -0700 Subject: [PATCH 11/14] Rename input variable to be nicer with workload initializer --- python/tvm/relay/testing/lstm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/lstm.py b/python/tvm/relay/testing/lstm.py index bc3991cb2434..47e68a988dab 100644 --- a/python/tvm/relay/testing/lstm.py +++ b/python/tvm/relay/testing/lstm.py @@ -133,7 +133,7 @@ def get_net(iterations, num_hidden, batch_size=1, dtype="float32"): out = None for i in range(iterations): - inputs = relay.Var("inputs_%s" % i, input_type) + inputs = relay.Var("data", input_type) i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type) i2h_bias = relay.Var("i2h_%i_bias" % i, bias_type) h2h_weight = relay.Var("h2h_%s_weight" % i, weight_type) From 2d4643d9122d67f61d7889830dac40e4bc71c2c4 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Fri, 2 Nov 2018 18:46:28 -0700 Subject: [PATCH 12/14] Split should not reject axis of 0 --- src/relay/op/tensor/transform.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 20e0e3adbfd3..eb9d778c66ca 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -910,7 +910,7 @@ bool SplitRel(const Array& types, } CHECK_LT(axis, data->shape.size()) << "axis should be within the input dimension range."; - CHECK_GT(axis, 0) + CHECK_GE(axis, 0) << "axis should be within the input dimension range."; if (const IntImm* sections = param->indices_or_sections.as()) { From c85be8ecf2009a1831518779b17ffe63ed7b72e7 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Fri, 2 Nov 2018 18:48:01 -0700 Subject: [PATCH 13/14] Add axis = 0 regression test for split --- tests/python/relay/test_op_level3.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index 26eccf991d0e..9aa75e77e8ce 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -160,6 +160,14 @@ def verify_split(dshape, indices_or_sections, ret_type, axis=None): relay.ty.TensorType((5, 1, 2, 2), "float32"), relay.ty.TensorType((5, 1, 2, 2), "float32")])), axis=1) + verify_split((5, 5, 2, 2), 5, + relay.ty.TupleType(tvm.convert([ + relay.ty.TensorType((1, 5, 2, 2), "float32"), + relay.ty.TensorType((1, 5, 2, 2), "float32"), + relay.ty.TensorType((1, 5, 2, 2), "float32"), + relay.ty.TensorType((1, 5, 2, 2), "float32"), + relay.ty.TensorType((1, 5, 2, 2), "float32")])), + axis=0) verify_split((d1, d2, d3, d4), 4, relay.ty.TupleType(tvm.convert([ relay.ty.TensorType((d1, d2, d3/4, d4), "float32"), From 2143e21aec6441b4fd0d486dec2aadb1bed81672 Mon Sep 17 00:00:00 2001 From: "Steven S. Lyubomirsky" Date: Fri, 2 Nov 2018 18:53:52 -0700 Subject: [PATCH 14/14] Symbolic regression test for split on axis zero --- tests/python/relay/test_op_level3.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index 9aa75e77e8ce..cd471e15d50b 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -175,6 +175,11 @@ def verify_split(dshape, indices_or_sections, ret_type, axis=None): relay.ty.TensorType((d1, d2, d3/4, d4), "float32"), relay.ty.TensorType((d1, d2, d3/4, d4), "float32")])), axis=2) + verify_split((d1, d2, d3, d4), 2, + relay.ty.TupleType(tvm.convert([ + relay.ty.TensorType((d1/2, d2, d3, d4), "float32"), + relay.ty.TensorType((d1/2, d2, d3, d4), "float32")])), + axis=0) verify_split((d1, d2, d3, d4), (2, 4, 7), relay.ty.TupleType(tvm.convert([ relay.ty.TensorType((d1, 2, d3, d4), "float32"),