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,76 @@
+# 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.convtranspose to align dynamic output_size input
+    """
+
+    def __init__(self):
+        """
+        Matches: functional batch_norm
+        """
+        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:
+            correct 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 functional.batch_norm and its inputs with BatchNorm2d.
+        Args:
+            graph: Graph we apply the substitution on.
+            node: node that match the pattern in the substitution init.
+
+        Returns:
+            Graph after applying the substitution.
+        """
+        # Check that output is only contain a single tensor
+        if len(node.output_shape) > 1:
+            Logger.critical('Output to nn.ConvTranspose2d should be a single tensor but got more than one.')  # pragma: no cover
+        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/pytorch_tests/model_tests/feature_models/convtranspose_dynamic_output_size_test.py

@@ -0,0 +1,58 @@
+# 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
+import torch.nn as nn
+from tests.pytorch_tests.model_tests.base_pytorch_test import BasePytorchTest
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+
+"""
+This test checks dynamic output size for nn.ConvTranspose2d.
+"""
+
+
+class ConvTranspose2dDynamicNet(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)
+
+    def forward(self, x):
+        x = self.downsample(x)
+        x = self.upsample(x, output_size=[224, 224])  # <--- dynamic output_size
+        return x
+
+
+class ConvTranspose2dDynamicNetTest(BasePytorchTest):
+    """
+    This test checks the addition and subtraction operations.
+    Both with different layers and with constants.
+    """
+
+    def __init__(self, unit_test):
+        super().__init__(unit_test)
+
+    def create_inputs_shape(self):
+        return [[self.val_batch_size, 3, 224, 224]]
+
+    def compare(self, quantized_models, float_model, input_x=None, quantization_info=None):
+        in_torch_tensor = to_torch_tensor(input_x[0])
+        for _, qmodel in quantized_models.items():
+            y_float = float_model(in_torch_tensor)
+            y_quant = qmodel(in_torch_tensor)
+            self.unit_test.assertTrue(y_float.shape == y_quant.shape,
+                                      msg=f'Out shape of the quantized model is not as the float model!')
+
+    def create_feature_network(self, input_shape):
+        return ConvTranspose2dDynamicNet()

tests/pytorch_tests/model_tests/test_feature_models_runner.py

@@ -47,6 +47,8 @@
 from tests.pytorch_tests.model_tests.feature_models.constant_conv_substitution_test import ConstantConvSubstitutionTest, \
     ConstantConvReuseSubstitutionTest, ConstantConvTransposeSubstitutionTest
 from tests.pytorch_tests.model_tests.feature_models.conv2d_replacement_test import DwConv2dReplacementTest
+from tests.pytorch_tests.model_tests.feature_models.convtranspose_dynamic_output_size_test import \
+    ConvTranspose2dDynamicNetTest
 from tests.pytorch_tests.model_tests.feature_models.dynamic_size_inputs_test import ReshapeNetTest
 from tests.pytorch_tests.model_tests.feature_models.gptq_test import GPTQAccuracyTest, GPTQWeightsUpdateTest, \
     GPTQLearnRateZeroTest
@@ -877,7 +879,11 @@ def test_manual_bit_width_selection_by_layer_name(self):
         ManualBitWidthByLayerNameTest(self, [NodeNameFilter('add'), NodeNameFilter('conv1_bn')], [2, 4]).run_test()
         ManualBitWidthByLayerNameTest(self, [NodeNameFilter('add'), NodeNameFilter('conv1_bn')], 4).run_test()
 
-
+    def test_convtranspose_dynamic_output_size(self):
+        """
+        This tests checks nn.ConvTranspose2d substitution for dynamic output size
+        """
+        ConvTranspose2dDynamicNetTest(self).run_test()
 
 
 if __name__ == '__main__':