diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 397ca90de55f..75ebb44cbb84 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -55,6 +55,17 @@ def revert_caffe2_pad(pads): return pads +def get_pad_pair(input1d, kernel1d, stride1d): + """infer pad size""" + if input1d % stride1d == 0: + pad = max(kernel1d - stride1d, 0) + else: + pad = max(kernel1d - (input1d % stride1d), 0) + pad_before = pad // 2 + pad_after = pad - pad_before + return [pad_before, pad_after] + + def onnx_storage_order2layout(storage_order): """converter of onnx storage order parameter to tvm storage order format""" if storage_order not in (0, 1): @@ -182,14 +193,37 @@ class Conv(OnnxOpConverter): @classmethod def _impl_v1(cls, inputs, attr, params): - out = AttrCvt(op_name=dimension_picker('conv'), - transforms={ - 'kernel_shape': 'kernel_size', - 'dilations': ('dilation', (0, 0)), - 'pads': ('padding', (0, 0), revert_caffe2_pad), - 'group': ('groups', 1)}, - ignores=['auto_pad'], - custom_check=dimension_constraint())(inputs[:2], attr, params) + # infer pads for auto_pad + if 'auto_pad' in attr: + attr['auto_pad'] = attr['auto_pad'].decode('utf-8') + if attr['auto_pad'] in ('SAME_UPPER', 'SAME_LOWER'): + input_shape = infer_shape(inputs[0]) + in_h, in_w = input_shape[2], input_shape[3] + stride_h, stride_w = attr['strides'] + kernel_h, kernel_w = attr['kernel_shape'] + dilation_h, dilation_w = attr['dilations'] + dilated_kernel_h = (kernel_h - 1) * dilation_h + 1 + dilated_kernel_w = (kernel_w - 1) * dilation_w + 1 + pad_v = get_pad_pair(in_h, dilated_kernel_h, stride_h) + pad_h = get_pad_pair(in_w, dilated_kernel_w, stride_w) + attr['pads'] = (pad_v[0], pad_h[0], pad_v[1], pad_h[1]) + elif attr['auto_pad'] == 'VALID': + attr['pads'] = (0, 0) + elif attr['auto_pad'] == 'NOTSET': + pass + else: + msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' + raise tvm.error.OpAttributeInvalid(msg.format(attr['auto_pad'])) + attr.pop('auto_pad') + + out = AttrCvt( + op_name=dimension_picker('conv'), + transforms={ + 'kernel_shape': 'kernel_size', + 'dilations': ('dilation', (0, 0)), + 'pads': ('padding', (0, 0), revert_caffe2_pad), + 'group': ('groups', 1)}, + custom_check=dimension_constraint())(inputs[:2], attr, params) use_bias = len(inputs) == 3 if use_bias: out = _op.nn.bias_add(out, inputs[2]) @@ -206,6 +240,29 @@ def _impl_v1(cls, inputs, attr, params): attr['channels'] = channels groups = attr.pop('group') attr['groups'] = groups + # infer pads for auto_pad + if 'auto_pad' in attr: + attr['auto_pad'] = attr['auto_pad'].decode('utf-8') + if attr['auto_pad'] in ('SAME_UPPER', 'SAME_LOWER'): + input_shape = infer_shape(inputs[0]) + in_h, in_w = input_shape[2], input_shape[3] + stride_h, stride_w = attr['strides'] + kernel_h, kernel_w = attr['kernel_shape'] + dilation_h, dilation_w = attr['dilations'] + dilated_kernel_h = (kernel_h - 1) * dilation_h + 1 + dilated_kernel_w = (kernel_w - 1) * dilation_w + 1 + pad_v = get_pad_pair(in_h, dilated_kernel_h, stride_h) + pad_h = get_pad_pair(in_w, dilated_kernel_w, stride_w) + attr['pads'] = (pad_v[0], pad_h[0], pad_v[1], pad_h[1]) + elif attr['auto_pad'] == 'VALID': + attr['pads'] = (0, 0) + elif attr['auto_pad'] == 'NOTSET': + pass + else: + msg = 'Value {} in attribute "auto_pad" of operator Conv is invalid.' + raise tvm.error.OpAttributeInvalid(msg.format(attr['auto_pad'])) + attr.pop('auto_pad') + out = AttrCvt( op_name=dimension_picker('conv', '_transpose'), transforms={