Support output_padding in ConvTranspose while doing ONNX exporting (#4583)
diff --git a/torch/onnx/symbolic.py b/torch/onnx/symbolic.py
index c59924e..c6583ba 100644
--- a/torch/onnx/symbolic.py
+++ b/torch/onnx/symbolic.py
@@ -376,9 +376,6 @@
def _convolution(g, input, weight, bias, stride, padding, dilation,
transposed, output_padding, groups, benchmark, deterministic, cudnn_enabled):
- if any(o != 0 for o in output_padding):
- return _unimplemented("_convolution", "non-zero output_padding")
-
weight_size = weight.type().sizes()
args = [input, weight]
@@ -394,6 +391,14 @@
"dilations_i": dilation,
"group_i": groups}
+ if any(o != 0 for o in output_padding):
+ # ONNX supports both output_shape and output_padding. they are equivalent expressive.
+ # output_padding is more straightforward, so we use it here.
+ # output_shape = stride * (input_shape - 1) + output_padding + kernel_shape - padding * 2
+ assert transposed
+ assert len(stride) == len(output_padding)
+ kwargs["output_padding_i"] = output_padding
+
n = g.op("ConvTranspose" if transposed else "Conv", *args, **kwargs)
if bias.node().kind() != "Undefined" and len(bias.type().sizes()) != 1:
