from __future__ import absolute_import
import topi
-from topi.util import get_const_int, get_const_tuple
+from topi.util import get_const_tuple
from .. import op as reg
from ..op import OpPattern, schedule_injective
if dilation_h < 1 or dilation_w < 1:
raise ValueError("dilation should be positive value")
+ def _get_out_depth():
+ weight_shape = get_const_tuple(inputs[1].shape)
+ if kernel_layout == "HWOI":
+ return weight_shape[2] * weight_shape[3]
+ return weight_shape[0] * weight_shape[1]
+
if groups == 1:
out = topi.nn.conv2d(
inputs[0], inputs[1], strides, padding,
dilation, layout, out_dtype)
- elif layout == "NCHW" and \
- get_const_int(inputs[1].shape[0]) == groups and \
- get_const_int(inputs[1].shape[1]) == 1:
+ elif layout == "NCHW" and _get_out_depth() == groups:
out = topi.nn.depthwise_conv2d_nchw(
inputs[0], inputs[1], strides, padding, dilation, out_dtype)
- elif layout == "NHWC" and \
- kernel_layout == "HWOI" and\
- get_const_int(inputs[1].shape[2]) == groups and \
- get_const_int(inputs[1].shape[3]) == 1:
+ elif layout == "NHWC" and kernel_layout == "HWOI" and _get_out_depth() == groups:
out = topi.nn.depthwise_conv2d_nhwc(
inputs[0], inputs[1], strides, padding, dilation, out_dtype)
elif layout in ['NCHW', 'NCHW4c']:
* \brief Convolution operators
*/
#include <tvm/data_layout.h>
+#include <tvm/ir_pass.h>
#include <tvm/relay/op.h>
#include <tvm/relay/attrs/nn.h>
#include <vector>
if (param->kernel_size.defined() && param->channels.defined()) {
CHECK_EQ(param->kernel_size.size(), 2);
CHECK_EQ(param->dilation.size(), 2);
- Array<IndexExpr> wshape(
- {param->channels,
- dshape_nchw[1] / param->groups,
- param->kernel_size[0],
- param->kernel_size[1]});
+ Array<IndexExpr> wshape;
+
+ if (tvm::ir::Equal(param->channels, param->groups)) {
+ // infer weight's shape for depthwise convolution
+ wshape = {
+ {dshape_nchw[1],
+ param->groups / dshape_nchw[1],
+ param->kernel_size[0],
+ param->kernel_size[1]}};
+ } else {
+ wshape = {
+ {param->channels,
+ dshape_nchw[1] / param->groups,
+ param->kernel_size[0],
+ param->kernel_size[1]}};
+ }
+
wshape = trans_kernel_layout.BackwardShape(wshape);
channels = param->channels;
dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
_test_convolution('depthwise', [4, 19, 17, 17], [3, 3, 19, 1], [1, 1], [2, 2], 'VALID', 'NCHW')
_test_convolution('depthwise', [4, 124, 17, 17], [1, 1, 124, 1], [1, 1], [1, 1], 'SAME', 'NCHW')
_test_convolution('depthwise', [4, 12, 17, 17], [3, 3, 12, 1], [1, 1], [2, 2], 'VALID', 'NCHW')
+ _test_convolution('depthwise', [4, 12, 17, 17], [3, 3, 12, 2], [1, 1], [2, 2], 'VALID', 'NCHW')
_test_convolution('conv', [4, 8, 8, 176], [1, 1, 176, 32], [1, 1], [1, 1], 'SAME', 'NHWC')
_test_convolution('conv', [4, 17, 17, 19], [3, 3, 19, 19], [1, 1], [2, 2], 'VALID', 'NHWC')
_test_convolution('depthwise', [4, 17, 17, 19], [3, 3, 19, 1], [1, 1], [2, 2], 'VALID', 'NHWC')
_test_convolution('depthwise', [4, 17, 17, 124], [1, 1, 124, 1], [1, 1], [1, 1], 'SAME', 'NHWC')
_test_convolution('depthwise', [4, 17, 17, 12], [3, 3, 12, 1], [1, 1], [2, 2], 'VALID', 'NHWC')
+ _test_convolution('depthwise', [4, 17, 17, 12], [3, 3, 12, 2], [1, 1], [2, 2], 'VALID', 'NHWC')
#######################################################################
# BiasAdd