add dynamic padding substitution for nn.ConvTranspose2d (#1381)

* add dynamic padding substitution for nn.ConvTranspose2d

Author: eladc-git

model_compression_toolkit/core/pytorch/constants.py

@@ -33,6 +33,7 @@
 DILATIONS = 'dilation'
 TENSOR_META = 'tensor_meta'
 FILTERS = 'out_channels'
+OUTPUT_PADDING = 'output_padding'
 TYPE = 'type'
 PAD = 'pad'
 VALUE = 'value'

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/convtranspose_dynamic_padding.py

@@ -0,0 +1,77 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Tuple
+import torch.nn as nn
+import torch
+from model_compression_toolkit.core.pytorch.constants import OUTPUT_PADDING
+from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
+from model_compression_toolkit.core import common
+from model_compression_toolkit.core.common import BaseNode, Graph
+from model_compression_toolkit.logger import Logger
+
+
+class ConvtransposeDynamicPadding(common.BaseSubstitution):
+    """
+    Replace output_padding of nn.ConvTranspose2d to align dynamic output_size input.
+    In case there is a dynamic output_size in ConvTranspose2d forward function, we recalculate the
+    output_padding here according to node.output_shape (which is equal to the dynamic output_size if existed).
+    """
+
+    def __init__(self):
+        """
+        Matches: nn.ConvTranspose2d
+        """
+        convtr_node = NodeOperationMatcher(nn.ConvTranspose2d)
+        super().__init__(matcher_instance=convtr_node)
+
+
+    def calc_dynamic_output_size(self, node: BaseNode) -> Tuple[int]:
+        """
+        Calc the output padding to support dunamic output_size of nn.ConvTranspose2d
+        Args:
+            node: node to calculate output padding
+
+        Returns:
+            corrected output padding
+        """
+        convtr = nn.ConvTranspose2d(**node.framework_attr)
+        num_spatial_dims = 2
+        output_padding = convtr._output_padding(torch.randn(size=node.input_shape[0]),
+                                                node.output_shape[0],
+                                                convtr.stride,
+                                                convtr.padding,
+                                                convtr.kernel_size,
+                                                num_spatial_dims,
+                                                convtr.dilation)
+        return tuple(output_padding)
+
+
+    def substitute(self,
+                   graph: Graph,
+                   node: BaseNode) -> Graph:
+        """
+        Substitute nn.ConvTranspose2d with corrected output_padding for cases of dynamic output_size
+        Args:
+            graph: Graph we apply the substitution on.
+            node: node that match the pattern in the substitution init.
+
+        Returns:
+            Graph after applying the substitution.
+        """
+
+        if not node.reuse:
+            output_padding = self.calc_dynamic_output_size(node)
+            node.framework_attr.update({OUTPUT_PADDING: output_padding})
+        return graph

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -62,6 +62,8 @@
     ScaledDotProductDecomposition
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.transform_function_call_method import \
     TransformFunctionCallMethod
+from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.convtranspose_dynamic_padding import \
+    ConvtransposeDynamicPadding
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.const_holder_conv import \
     FunctionalConvSubstitution
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.relu_bound_to_power_of_2 import \
@@ -286,7 +288,8 @@ def get_substitutions_prepare_graph(self, fw_info: FrameworkInfo = None) -> List
                 FunctionalBatchNorm(),
                 FunctionalLayerNorm(),
                 FunctionalLinear(),
-                RemoveIdentity()]
+                RemoveIdentity(),
+                ConvtransposeDynamicPadding()]
 
     def get_substitutions_pre_statistics_collection(self,
                                                     quant_config: QuantizationConfig

tests_pytest/pytorch_tests/integration_tests/core/graph_substitutions/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tests_pytest/pytorch_tests/integration_tests/core/graph_substitutions/substitutions/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tests_pytest/pytorch_tests/integration_tests/core/graph_substitutions/substitutions/test_convtasnet_dynamic_output_size.py

@@ -0,0 +1,67 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+from torch import nn
+from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2pytorch import \
+    AttachTpcToPytorch
+
+from model_compression_toolkit.core import QuantizationConfig
+from model_compression_toolkit.core.graph_prep_runner import graph_preparation_runner
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.downsample = nn.Conv2d(3, 16, 3, stride=2, padding=1)
+        self.upsample = nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1)
+        self.downsample2 = nn.Conv2d(3, 16, 3, stride=2, padding=1)
+        self.upsample2 = nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1)
+
+    def forward(self, x):
+        x = self.downsample(x)
+        x = self.upsample(x, output_size=[224, 224])  # <--- dynamic output_size
+        x = self.downsample2(x)
+        x = self.upsample2(x)  # <--- no dynamic output_size
+        return x
+
+
+def data_gen():
+    yield [torch.rand(1, 3, 224, 224)]
+
+
+def test_convtranspose_dynamic_output_size(minimal_tpc):
+    Model()(next(data_gen())[0])
+
+    fw_impl = PytorchImplementation()
+    fw_info = DEFAULT_PYTORCH_INFO
+    model = Model()
+
+    graph = graph_preparation_runner(model,
+                                     data_gen,
+                                     QuantizationConfig(),
+                                     fw_info=fw_info,
+                                     fw_impl=fw_impl,
+                                     fqc=AttachTpcToPytorch().attach(minimal_tpc),
+                                     mixed_precision_enable=False,
+                                     running_gptq=False)
+
+    nodes = graph.get_topo_sorted_nodes()
+
+    assert nodes[2].framework_attr['output_padding'] == (1,1)
+    assert nodes[2].output_shape[0] == [1, 3, 224, 224]
+    assert nodes[4].framework_attr['output_padding'] == (0,0)
+    assert nodes[4].output_shape[0] == [1, 3, 223, 223]