[quantization] Enable nn.ConvTranspose2D in GPTQs (#428)

This PR enables quantization of `nn.ConvTranspose2D` in FPIGPTQ/GPTQ and adds tests for it.

TICO-DCO-1.0-Signed-off-by: s.malakhov <s.malakhov@partner.samsung.com>

Author: stamalakhov

test/quantization/algorithm/test_fpi_gptq.py

@@ -119,6 +119,24 @@ def get_example_inputs(self):
         return (torch.randn(1, 32, 16),), {}
 
 
+class TransposedConv2DGeneral(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.tconv = torch.nn.ConvTranspose2d(16, 32, (2, 2), stride=2, groups=1)
+        self.tconv2 = torch.nn.ConvTranspose2d(
+            32, 16, (3, 3), stride=4, groups=2
+        )  # general groupwise
+
+    def forward(self, x):
+        z = self.tconv(x)
+        z = self.tconv2(z)
+        return z
+
+    def get_example_inputs(self):
+        return (torch.randn(1, 16, 7, 7),), {}
+
+
 class FPIGPTQTest(unittest.TestCase):
     @unittest.skipIf(
         not IS_INTERNAL_TEST, "Internal test — run only if --include-internal is set"
@@ -321,3 +339,29 @@ def test_groupwise_conv1d(self):
         ), "second conv node is not quantized"
 
         # TODO add PT2E quantization (right now it can't be evaluated on backend)
+
+    @unittest.skipIf(
+        not IS_INTERNAL_TEST, "Internal test — run only if --include-internal is set"
+    )
+    def test_transposed_conv2d(self):
+        q_m = TransposedConv2DGeneral()
+        q_m.eval()
+        ori_m = q_m
+        args, kwargs = ori_m.get_example_inputs()
+
+        # Apply GPTQ
+        q_m = prepare(q_m, FPIGPTQConfig(show_progress=False))
+        for _ in range(30):
+            args, kwargs = ori_m.get_example_inputs()
+            q_m(*args, **kwargs)
+        convert(q_m, inplace=True)
+        # check that all convolution nodes are quantized
+        assert hasattr(q_m, "quantizers"), "quantized model does not have quantizers"
+        assert (
+            "model.layers.0.tconv" in q_m.quantizers
+        ), "first conv node is not quantized"
+        assert (
+            "model.layers.0.tconv2" in q_m.quantizers
+        ), "second conv node is not quantized"
+
+        # TODO add PT2E quantization

test/quantization/algorithm/test_gptq.py

@@ -119,6 +119,24 @@ def get_example_inputs(self):
         return (torch.randn(1, 32, 16),), {}
 
 
+class TransposedConv2DGeneral(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.tconv = torch.nn.ConvTranspose2d(16, 32, (2, 2), stride=2, groups=1)
+        self.tconv2 = torch.nn.ConvTranspose2d(
+            32, 16, (3, 3), stride=4, groups=2
+        )  # general groupwise
+
+    def forward(self, x):
+        z = self.tconv(x)
+        z = self.tconv2(z)
+        return z
+
+    def get_example_inputs(self):
+        return (torch.randn(1, 16, 7, 7),), {}
+
+
 class GPTQTest(unittest.TestCase):
     @unittest.skipIf(
         not IS_INTERNAL_TEST, "Internal test — run only if --include-internal is set"
@@ -352,3 +370,29 @@ def test_groupwise_conv1d(self):
         ), "second conv node is not quantized"
 
         # TODO add PT2E quantization (right now it can't be evaluated on backend)
+
+    @unittest.skipIf(
+        not IS_INTERNAL_TEST, "Internal test — run only if --include-internal is set"
+    )
+    def test_transposed_conv2d(self):
+        q_m = TransposedConv2DGeneral()
+        q_m.eval()
+        ori_m = q_m
+        args, kwargs = ori_m.get_example_inputs()
+
+        # Apply GPTQ
+        q_m = prepare(q_m, GPTQConfig(show_progress=False))
+        for _ in range(30):
+            args, kwargs = ori_m.get_example_inputs()
+            q_m(*args, **kwargs)
+        convert(q_m, inplace=True)
+        # check that all convolution nodes are quantized
+        assert hasattr(q_m, "quantizers"), "quantized model does not have quantizers"
+        assert (
+            "model.layers.0.tconv" in q_m.quantizers
+        ), "first conv node is not quantized"
+        assert (
+            "model.layers.0.tconv2" in q_m.quantizers
+        ), "second conv node is not quantized"
+
+        # TODO add PT2E quantization

tico/quantization/algorithm/fpi_gptq/fpi_gptq.py

@@ -25,6 +25,12 @@
 import torch
 import torch.nn as nn
 
+from tico.quantization.algorithm.gptq.gptq import (
+    conv2d_weights_to_convtranspose2d_weights,
+    convtranspose2d_weights_to_conv2d_weights,
+    get_matmul_input_for_convtranspose2d,
+)
+
 from tico.quantization.algorithm.gptq.quant import quantize, Quantizer
 
 
@@ -58,6 +64,9 @@ def __init__(self, layer):
         W = layer.weight.data.clone()
         if isinstance(self.layer, nn.Conv2d) or isinstance(self.layer, nn.Conv1d):
             W = W.flatten(1)
+        elif isinstance(self.layer, nn.ConvTranspose2d):
+            W = convtranspose2d_weights_to_conv2d_weights(self.layer, W)
+            W = W.flatten(1)
 
         self.rows = W.shape[0]
         self.columns = W.shape[1]
@@ -132,6 +141,8 @@ def add_batch(self, inp, out):
             inp = unfold(inp)
             inp = inp.permute([1, 0, 2])
             inp = inp.flatten(1)
+        if isinstance(self.layer, nn.ConvTranspose2d):
+            inp = get_matmul_input_for_convtranspose2d(self.layer, inp)
 
         self.H *= self.nsamples / (self.nsamples + tmp)
         self.nsamples += tmp
@@ -146,6 +157,11 @@ def fasterquant(
         W = self.layer.weight.data.clone()
         if isinstance(self.layer, nn.Conv2d) or isinstance(self.layer, nn.Conv1d):
             W = W.flatten(1)
+        elif isinstance(self.layer, nn.ConvTranspose2d):
+            W = convtranspose2d_weights_to_conv2d_weights(self.layer, W)
+            conv2d_shape = W.shape
+            W = W.flatten(1)  # reshaped to matrix (OUT_channels x the_rest)
+
         W = W.float()
         tick = time.time()
         if not self.quantizer.ready():
@@ -202,6 +218,15 @@ def fasterquant(
                 self.quantizer.zero,
                 self.quantizer.maxq,
             )
+        elif isinstance(self.layer, nn.ConvTranspose2d):
+            Q[:, dead] = quantize(
+                convtranspose2d_weights_to_conv2d_weights(
+                    self.layer, self.layer.weight.data
+                ).flatten(1)[:, dead],
+                self.quantizer.scale,
+                self.quantizer.zero,
+                self.quantizer.maxq,
+            )
         else:
             Q[:, dead] = quantize(
                 self.layer.weight[:, dead],
@@ -210,9 +235,15 @@ def fasterquant(
                 self.quantizer.maxq,
             )
 
-        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(
-            self.layer.weight.data.dtype
-        )
+        if isinstance(self.layer, nn.ConvTranspose2d):
+            Q_conv2d = Q.reshape(conv2d_shape).to(self.layer.weight.data.dtype)
+            self.layer.weight.data = conv2d_weights_to_convtranspose2d_weights(
+                self.layer, Q_conv2d
+            )
+        else:
+            self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(
+                self.layer.weight.data.dtype
+            )
 
     def free(self):
         self.H = None

tico/quantization/algorithm/fpi_gptq/quantizer.py

@@ -77,7 +77,13 @@ def convert(self, model):
         ):
             # 1) Identify quantizable submodules within the layer
             full = find_layers(
-                layer, layers=[torch.nn.Linear, torch.nn.Conv2d, torch.nn.Conv1d]
+                layer,
+                layers=[
+                    torch.nn.Linear,
+                    torch.nn.Conv2d,
+                    torch.nn.Conv1d,
+                    torch.nn.ConvTranspose2d,
+                ],
             )
             sequential = [list(full.keys())]