Support Deriving channels when it is not provided in AlterLayout. (#2972)

tests/python/frontend/nnvm_to_relay/test_alter_conv2d.py

+# 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.
+"""Test alter conv2d layout pass"""
+import tvm
+import nnvm
+
+from tvm import relay
+from tvm import autotvm
+from tvm.relay.ir_pass import infer_type, alpha_equal
+
+
+def test_alter_layout_conv2d():
+    """Additional layout transformations should occour on the graph.
+    """
+
+    def convnet():
+        """Alternating layout of simple convnet (from image super-resolution).
+        """
+        bias1 = relay.var('bias1', shape=(64,))
+        bias2 = relay.var('bias2', shape=(64,))
+        bias3 = relay.var('bias3', shape=(64,))
+        bias4 = relay.var('bias4', shape=(64,))
+        weight1 = relay.var('weight1', shape=(64, 1, 5, 5))
+        weight2 = relay.var('weight2', shape=(64, 64, 3, 3))
+        weight3 = relay.var('weight3', shape=(64, 64, 3, 3))
+        weight4 = relay.var('weight4', shape=(64, 64, 3, 3))
+        data = relay.var("x", shape=(1, 1, 224, 224))
+        n00 = relay.nn.conv2d(data, weight1, padding=[2, 2], kernel_size=[5, 5])
+        n01 = relay.expand_dims(bias1, axis=1, num_newaxis=2)
+        n02 = relay.add(n00, n01)
+        n03 = relay.nn.relu(n02)
+        n04 = relay.nn.conv2d(n03, weight2, padding=[1, 1], kernel_size=[3, 3])
+        n05 = relay.expand_dims(bias2, axis=1, num_newaxis=2)
+        n06 = relay.add(n04, n05)
+        n07 = relay.nn.relu(n06)
+        n08 = relay.nn.conv2d(n07, weight3, padding=[1, 1], kernel_size=[3, 3])
+        n09 = relay.expand_dims(bias3, axis=1, num_newaxis=2)
+        n10 = relay.add(n08, n09)
+        n11 = relay.nn.relu(n10)
+        n12 = relay.nn.conv2d(n11, weight4, padding=[1, 1], kernel_size=[3, 3])
+        n13 = relay.expand_dims(bias4, axis=1, num_newaxis=2)
+        n14 = relay.add(n12, n13)
+        n15 = relay.reshape(n14, newshape=[1, 1, 3, 3, 224, 224])
+        n16 = relay.transpose(n15, axes=[0, 1, 4, 2, 5, 3])
+        net = relay.reshape(n16, newshape=[1, 1, 672, 672])
+        args = relay.ir_pass.free_vars(net)
+        return relay.Function(args, net)
+
+    # orig net
+    N = convnet()
+    N = infer_type(N)
+
+    # trigger a test
+    # for each known alter_conv2d
+    targets=['cuda',
+             'opencl -device=mali',
+             'opencl -device=intel_graphics',
+             'llvm -device=arm_cpu',
+             'llvm -device=core-avx-ii']
+
+    for tgt in targets:
+        with tvm.target.create(tgt) as target:
+            with relay.build_config(opt_level=-1, add_pass='AlterOpLayout'):
+               with autotvm.tophub.context(target):
+                   O = relay.optimize(N, target, params=None)
+                   O = relay.ir_pass.infer_type(O)
+
+                   # graph should differ
+                   assert not relay.ir_pass.alpha_equal(N, O)
+
+if __name__ == "__main__":
+    np.random.seed(42)
+    test_alter_layout_conv2d()

topi/python/topi/arm_cpu/conv2d.py

@@ -700,6 +700,10 @@ def _alter_conv2d_layout_arm(attrs, inputs, tinfos, F):
 
     new_attrs = {k: attrs[k] for k in attrs.keys()}
 
+    if F == tvm.relay.op:
+        # Derive channels for frontends (e.g ONNX) that miss "channel" field.
+        new_attrs["channels"] = inputs[1].checked_type.shape[attrs['kernel_layout'].index('O')]
+
     dilation = attrs.get_int_tuple("dilation")
     strides = attrs.get_int_tuple("strides")
     padding = attrs.get_int_tuple("padding")

topi/python/topi/cuda/conv2d_winograd.py

@@ -371,6 +371,10 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, F):
     copy_inputs = [s for s in inputs]
     new_attrs = {k: attrs[k] for k in attrs.keys()}
 
+    if F == tvm.relay.op:
+        # Derive channels for frontends (e.g ONNX) that miss "channel" field.
+        new_attrs["channels"] = inputs[1].checked_type.shape[attrs['kernel_layout'].index('O')]
+
     strides = attrs.get_int_tuple("strides")
     padding = attrs.get_int_tuple("padding")
     dilation = attrs.get_int_tuple("dilation")

topi/python/topi/intel_graphics/conv2d.py

@@ -73,7 +73,11 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, F):
             break
 
     new_attrs = {k: attrs[k] for k in attrs.keys()}
-    new_attrs['kernel_layout'] = 'OIHW%do' % (oc_bn)
+    new_attrs["kernel_layout"] = 'OIHW%do' % (oc_bn)
+
+    if F == tvm.relay.op:
+        # Derive channels for frontends (e.g ONNX) that miss "channel" field.
+        new_attrs["channels"] = inputs[1].checked_type.shape[attrs['kernel_layout'].index('O')]
 
     if F == sym:
         out = F.contrib.conv2d_NCHWc(*copy_inputs, **new_attrs)

topi/python/topi/x86/conv2d.py

@@ -327,11 +327,16 @@ def _alter_conv2d_layout(attrs, inputs, tinfo, F):
 
     copy_inputs = [s for s in inputs]
     new_attrs = {k : attrs[k] for k in attrs.keys()}
+
+    if F == tvm.relay.op:
+        # Derive channels for frontends (e.g ONNX) that miss "channel" field.
+        new_attrs["channels"] = inputs[1].checked_type.shape[attrs['kernel_layout'].index('O')]
+
     data, kernel = tinfo[0], tinfo[1]
     batch_size, in_channel, height, width = get_const_tuple(data.shape)
 
     groups = attrs.get_int("groups")
-    out_channel = attrs.get_int("channels") if F == sym else attrs.get_int("channels").value
+    out_channel = attrs.get_int("channels") if F == sym else new_attrs["channels"]
     padding = attrs.get_int_tuple("padding")
     strides = attrs.get_int_tuple("strides")
     dilation = attrs.get_int_tuple("dilation")