Mismatching between schedule and compute for group conv2d

[kumasento]

Hello,

I'm not sure whether anyone has met this problem before, but when I was trying to run an ONNX model that has group convolution, i.e., group > 1 && group != in_channels, Relay will apply mismatched schedule and compute function (see here):

# schedule
topi.generic.schedule_depthwise_conv2d_nchw(outs)
# compute
topi.nn.group_conv2d_nchw(inputs[0], inputs[1], strides, padding, dilation, groups,
                                              out_dtype=out_dtype)

The reason behind is that some shape-based conditions for making the compute function decision are not available for scheduling.

I tried to solve this issue, but based on the input arguments that schedule_conv2d receives, i.e., attrs, outs, I cannot infer whether a given conv2d op should be dispatched to depthwise_conv2d or group_conv2d. At least we need to have the number of input channels to properly decide which convolution op to go with.

Therefore, I intend to extend Conv2DAttrs a little bit with an extra field called in_channels (at here), and then update the schedule_conv2d function to take it into account.

Please let me know whether I've missed anything. Thanks!

[kumasento]

To reproduce this issue, please check this short example:
gconv_tvm_onnx.zip

# requires onnx
python create_onnx_model.py
python load_gconv.py

I'm using the latest master of TVM (4335db0)