[Optimization] Apply multi times apply passes will result in check shape error

This is the reduced case from real model leveraging test_fold_scale_axis.py

import nnvm
import nnvm.testing.resnet
import numpy as np
from nnvm import symbol as sym
from nnvm.compiler import graph_util, graph_attr

def test_fold_axis_conv():
    def before(x, conv_weight, conv_bias, in_scale, out_scale, channels):
        x = x * sym.expand_dims(in_scale, axis=1, num_newaxis=2)
        y = sym.conv2d(x, conv_weight, conv_bias,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       groups=54,
                       name="conv")
        y = sym.relu(y)
        y = y * sym.expand_dims(out_scale, axis=1, num_newaxis=2)
        return y

    def expected(x, conv_weight, conv_bias, in_scale, out_scale, channels):
        conv_weight = conv_weight * sym.expand_dims(out_scale, axis=1, num_newaxis=3)
        conv_weight = conv_weight * sym.expand_dims(in_scale, axis=1, num_newaxis=2)
        conv_bias = conv_bias * out_scale
        y = sym.conv2d(x,
                       conv_weight,
                       conv_bias,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       groups=54,
                       name="conv")
        y = sym.relu(y)
        return y

    # Before simplify
    def check(shape, channels):
        x = sym.Variable("x") + 1
        weight = sym.Variable("weight", shape=(54, 1, 3, 3))
        bias = sym.Variable("bias", shape=(54,))
        in_scale = sym.Variable("in_scale")
        out_scale = sym.Variable("out_scale")
        y1 = before(x, weight, bias, in_scale, out_scale, channels)
        y2 = expected(x, weight, bias, in_scale, out_scale, channels)
        ishape = {"x": shape, "out_scale": (channels,), "in_scale": (shape[1],)}
        g1 = nnvm.graph.create(y1)
        g2 = nnvm.graph.create(y2)
        graph_attr.set_shape_inputs(g1, ishape)
        g1 = g1.apply("InferShape").apply("FoldScaleAxis")
        # The second time
        graph_attr.set_shape_inputs(g1, ishape)
        g1 = g1.apply("InferShape").apply("FoldScaleAxis")

    check((1, 54, 63, 27), 54)
test_fold_axis_conv()

The problem is the second time. Will result in:

Traceback (most recent call last): File “/Users/blue/Documents/dev/tmp/tvm/nnvm/python/nnvm/_base.py”, line 75, in check_call raise NNVMError(py_str(_LIB.NNGetLastError())) nnvm._base.NNVMError: [19:28:33] /Users/blue/Documents/dev/tmp/tvm/nnvm/src/top/nn/convolution.cc:91: Operator conv2d(padding=(1, 1), kernel_size=(3, 3), channels=54, name=conv) expects weight’s shape to be [54,54,3,3], but got [54,1,3,3].

The second time to apply optimization is the same as apply it into g2.

[FrozenGene]

@merrymercy I am applying the process of open source in my company, could you wait for a while? Sorry for this inconvenience.

[FrozenGene]

The reason is FoldScaleAxis doesn’t support convolution with groups (including depthwise conv2d) currently. The fix can be done modify ScaleAxisForward for conv2d. Close it.