[TESTS] Decrease test times by introducing testing model (#6235)

Adds a new testing model `tvm.relay.testing.synthetic` which is a small,
but representative model. Replaces resnet with this model in many tests.

python/tvm/relay/testing/__init__.py

@@ -39,6 +39,7 @@
 from . import densenet
 from . import yolo_detection
 from . import temp_op_attr
+from . import synthetic
 
 from .config import ctx_list
 from .init import create_workload

python/tvm/relay/testing/init.py

@@ -15,11 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 """Initializer of parameters."""
+from functools import reduce
 import numpy as np
 
 import tvm
 from tvm import relay
 
+
 class Initializer(object):
     """The base class of an initializer."""
     def __init__(self, **kwargs):
@@ -128,6 +130,14 @@ def _init_weight(self, name, arr):
             raise ValueError("Unknown random type")
 
 
+class Constant(Initializer):
+    """ Constant initialization of weights. Sum of weights in the matrix is 1.
+    """
+    def _init_weight(self, name, arr):
+        num_elements = reduce(lambda x, y: x*y, arr.shape)
+        arr[:] = 1./num_elements
+
+
 def create_workload(net, initializer=None, seed=0):
     """Helper function to create benchmark image classification workload.
 

python/tvm/relay/testing/synthetic.py

@@ -0,0 +1,120 @@
+# 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.
+"""
+Synthetic networks for testing purposes. Ideally, these networks are similar in
+structure to real world networks, but are much smaller in order to make testing
+faster.
+"""
+from __future__ import absolute_import
+from tvm import relay
+from .init import create_workload, Constant
+from . import layers
+
+
+def get_net(input_shape=(1, 3, 24, 12), dtype="float32", wtype=None):
+    """Get synthetic testing network.
+
+    Parameters
+    ----------
+    image_shape : tuple, optional
+        The input shape as (batch_size, channels, height, width).
+
+    dtype : str, optional
+        The data type for the input.
+
+    wtype : str, optional
+        The data type for weights. Defaults to `dtype`.
+
+    Returns
+    -------
+    net : relay.Function
+        The dataflow.
+    """
+    if wtype is None:
+        wtype = dtype
+    data = relay.var("data", shape=input_shape, dtype=dtype)
+    dense_shape = [-1, input_shape[3]]
+    dense = relay.nn.relu(
+        relay.nn.dense(
+            relay.reshape(data, dense_shape),
+            relay.var(
+                "dense_weight", shape=[input_shape[3], dense_shape[1]], dtype=wtype
+            ),
+        )
+    )
+    dense = relay.reshape_like(dense, data)
+    conv_shape = [input_shape[1], input_shape[1], 3, 3]
+    conv = relay.nn.softmax(
+        relay.nn.conv2d(
+            data,
+            relay.var("conv_weight", shape=conv_shape, dtype=wtype),
+            padding=1,
+            kernel_size=3,
+        )
+    )
+    added = relay.add(dense, conv)
+    biased = layers.batch_norm_infer(
+        relay.nn.bias_add(added, relay.var("bias", dtype=wtype)), name="batch_norm"
+    )
+    dense = relay.nn.relu(
+        relay.nn.dense(
+            relay.reshape(biased, dense_shape),
+            relay.var(
+                "dense2_weight", shape=[input_shape[3], dense_shape[1]], dtype=wtype
+            ),
+        )
+    )
+    dense = relay.reshape_like(dense, data)
+    conv = relay.nn.softmax(
+        relay.nn.conv2d(
+            biased,
+            relay.var("conv2_weight", shape=conv_shape, dtype=wtype),
+            padding=1,
+            kernel_size=3,
+        )
+    )
+    added = relay.add(dense, conv)
+    args = relay.analysis.free_vars(added)
+    return relay.Function(args, added)
+
+
+def get_workload(input_shape=(1, 3, 24, 12), dtype="float32", wtype=None):
+    """Get benchmark workload for the synthetic net.
+
+    Parameters
+    ----------
+    image_shape : tuple, optional
+        The input shape as (batch_size, channels, height, width).
+
+    dtype : str, optional
+        The data type for the input.
+
+    wtype : str, optional
+        The data type for weights. Defaults to `dtype`.
+
+    Returns
+    -------
+    mod : tvm.IRModule
+        The relay module that contains a synthetic network.
+
+    params : dict of str to NDArray
+        The parameters.
+    """
+    return create_workload(
+        get_net(input_shape=input_shape, dtype=dtype, wtype=wtype),
+        initializer=Constant(),
+    )

tests/micro/test_runtime_micro_on_arm.py

@@ -23,7 +23,6 @@
 from tvm import relay
 import tvm.micro as micro
 from tvm.micro import create_micro_mod
-from tvm.relay.testing import resnet
 
 # Use real micro device - an STM32F746 discovery board
 # SETUP:

tests/python/relay/test_analysis_extract_fused_functions.py

@@ -17,7 +17,7 @@
 """Test function extraction"""
 import tvm
 from tvm import relay
-from tvm.relay.testing.resnet import get_workload
+from tvm.relay.testing.synthetic import get_workload
 
 
 def get_conv_net():
@@ -106,7 +106,7 @@ def is_conv_add(func):
 def test_extract_resnet():
     mod, _params = get_workload()
     items = relay.analysis.extract_fused_functions(mod)
-    assert len(items) == 34
+    assert len(items) == 6
 
 
 if __name__ == '__main__':

tests/python/relay/test_change_batch.py

@@ -17,13 +17,13 @@
 import tvm
 from tvm import te
 from tvm import relay
-from tvm.relay.testing import resnet
+from tvm.relay.testing import synthetic
 from tvm.relay import transform
 
-def test_change_batch_resnet():
-    net, params = resnet.get_workload()
+def test_change_batch_synthetic():
+    net, params = synthetic.get_workload()
     new_net = transform.ChangeBatch({net["main"].params[0]: 0}, batch_size=123)(net)
-    assert new_net["main"].checked_type.ret_type == relay.TensorType((123, 1000))
+    assert new_net["main"].checked_type.ret_type.shape[0] == 123
 
 if __name__ == "__main__":
-    test_change_batch_resnet()
+    test_change_batch_synthetic()

tests/python/relay/test_pass_auto_quantize.py

@@ -72,18 +72,19 @@ def _check_batch_flatten(node):
     # check if batch_flatten is quantized
     relay.analysis.post_order_visit(qmod["main"], _check_batch_flatten)
 
-def get_calibration_dataset(input_name):
+def get_calibration_dataset(mod, input_name):
     dataset = []
+    input_shape = [int(x) for x in mod["main"].checked_type.arg_types[0].shape]
     for i in range(5):
-        data = np.random.uniform(size=(1, 3, 224, 224))
+        data = np.random.uniform(size=input_shape)
         dataset.append({input_name: data})
     return dataset
 
 
 @pytest.mark.parametrize("create_target", [True, False])
 def test_calibrate_target(create_target):
-    mod, params = testing.resnet.get_workload(num_layers=18)
-    dataset = get_calibration_dataset("data")
+    mod, params = testing.synthetic.get_workload()
+    dataset = get_calibration_dataset(mod, "data")
     with relay.quantize.qconfig(calibrate_mode="kl_divergence"):
         if create_target:
             with tvm.target.create("llvm"):
@@ -94,8 +95,8 @@ def test_calibrate_target(create_target):
 
 
 def test_calibrate_memory_bound():
-    mod, params = testing.resnet.get_workload(num_layers=18)
-    dataset = get_calibration_dataset("data")
+    mod, params = testing.synthetic.get_workload()
+    dataset = get_calibration_dataset(mod, "data")
     import multiprocessing
     num_cpu = multiprocessing.cpu_count()
     with relay.quantize.qconfig(calibrate_mode="kl_divergence",

tests/python/relay/test_vm.py

@@ -594,7 +594,7 @@ def test_add_op_broadcast():
     check_result([x_data, y_data], x_data + y_data, mod=mod)
 
 def test_vm_optimize():
-    mod, params = testing.resnet.get_workload(batch_size=1, num_layers=18)
+    mod, params = testing.synthetic.get_workload()
     comp = relay.vm.VMCompiler()
     opt_mod, _ = comp.optimize(mod, "llvm", params)
 

tests/python/relay/test_vm_serialization.py

@@ -51,13 +51,14 @@ def get_serialized_output(mod, *data, params=None, target="llvm",
 
 def run_network(mod,
                 params,
-                data_shape=(1, 3, 224, 224),
                 dtype='float32'):
     def get_vm_output(mod, data, params, target, ctx, dtype='float32'):
         ex = relay.create_executor('vm', mod=mod, ctx=ctx)
         result = ex.evaluate()(data, **params)
         return result.asnumpy().astype(dtype)
 
+    print(mod["main"])
+    data_shape = [int(x) for x in mod["main"].checked_type.arg_types[0].shape]
     data = np.random.uniform(size=data_shape).astype(dtype)
     target = "llvm"
     ctx = tvm.cpu(0)
@@ -272,8 +273,8 @@ def test_closure():
     tvm.testing.assert_allclose(res.asnumpy(), 3.0)
 
 
-def test_resnet():
-    mod, params = testing.resnet.get_workload(batch_size=1, num_layers=18)
+def test_synthetic():
+    mod, params = testing.synthetic.get_workload()
     run_network(mod, params)
 
 
@@ -306,6 +307,6 @@ def test_vm_shape_of():
     test_adt_list()
     test_adt_compose()
     test_closure()
-    test_resnet()
+    test_synthetic()
     test_mobilenet()
     test_vm_shape_of()

tests/python/unittest/test_autotvm_graph_tuner_utils.py

@@ -24,7 +24,7 @@
 from tvm import te
 
 from tvm import autotvm, relay
-from tvm.relay.testing import resnet
+from tvm.relay.testing import synthetic
 from tvm.autotvm.graph_tuner.utils import has_multiple_inputs, get_direct_ancestor, get_in_nodes, \
     get_out_nodes, expr2graph, bind_inputs
 from tvm.autotvm.graph_tuner._base import OPT_OUT_OP
@@ -56,7 +56,7 @@ def test_has_multiple_inputs():
 
 
 def test_expr2graph():
-    mod, _ = resnet.get_workload(num_layers=50, batch_size=1)
+    mod, _ = synthetic.get_workload()
     node_dict = {}
     node_list = []
     target_ops = [relay.op.get("nn.conv2d")]

tests/python/unittest/test_runtime_micro.py

@@ -23,7 +23,6 @@
 from tvm import relay
 import tvm.micro as micro
 from tvm.micro import create_micro_mod
-from tvm.relay.testing import resnet
 
 # # Use the host emulated micro device.
 DEV_CONFIG_A = micro.device.host.generate_config()