[feat] Add IntTransposedConv

Author: chrisr-12

compressai_vision/model_wrappers/detectron2.py

@@ -1,4 +1,3 @@
-# Copyright (c) 2022-2024, InterDigital Communications, Inc
 # All rights reserved.
 
 # Redistribution and use in source and binary forms, with or without
@@ -46,7 +45,7 @@
 from compressai_vision.registry import register_vision_model
 
 from .base_wrapper import BaseWrapper
-from .intconv2d import IntConv2d
+from .intconv2d import IntConv2d, IntTransposedConv2d
 
 __all__ = [
     "faster_rcnn_X_101_32x8d_FPN_3x",
@@ -189,9 +188,13 @@ def size_divisibility(self):
 
     def replace_conv2d_modules(self, module):
         for child_name, child_module in module.named_children():
-            if type(child_module).__name__ == "Conv2d":
-                int_conv2d = Conv2d(**child_module.__dict__)
-                int_conv2d.set_attributes(child_module)
+            if type(child_module).__name__ in ["Conv2d", "TransposedConv2d"]:
+                if type(child_module).__name__ == "Conv2d":
+                    int_module = Conv2d(**child_module.__dict__)
+                    int_module.set_attributes(child_module)
+                else:
+                    int_module = IntTransposedConv2d(**child_module.__dict__)
+                    int_module.set_attributes(child_module)
 
                 # Since regular list is used instead of ModuleList
                 if "fpn_lateral" in child_name or "fpn_output" in child_name:
@@ -201,12 +204,12 @@ def replace_conv2d_modules(self, module):
                     assert idx in [2, 3, 4, 5]
 
                     if "fpn_lateral" in child_name:
-                        module.lateral_convs[3 - (idx - 2)] = int_conv2d
+                        module.lateral_convs[3 - (idx - 2)] = int_module
                     else:
                         assert "fpn_output" in child_name
-                        module.output_convs[3 - (idx - 2)] = int_conv2d
+                        module.output_convs[3 - (idx - 2)] = int_module
 
-                setattr(module, child_name, int_conv2d)
+                setattr(module, child_name, int_module)
             else:
                 self.replace_conv2d_modules(child_module)
 

compressai_vision/model_wrappers/intconv2d.py

@@ -147,3 +147,97 @@ def conv2d(self, x: torch.Tensor):
             self.dilation,
             self.groups,
         )
+
+
+class IntTransposedConv2d(torch.nn.ConvTranspose2d):
+    def __init__(self, *args, **kwarg) -> None:
+        _nkwargs = copy.deepcopy(kwargs)
+
+        del _nkwargs["training"]
+        del _nkwargs["transposed"]
+        del _nkwargs["output_padding"]
+        for name in kwargs.keys():
+            if name.startswith("_"):
+                del _nkwargs[name]
+
+        super().__init__(*args, **_nkwargs)
+        self.initified_weight_mode = False
+
+    # prepare quantized weights
+    def quantize(self):
+        self.initified_weight_mode = True
+
+        if self.bias is None:
+            self.float_bias = torch.zeros(self.out_channels, device=self.weight.device)
+        else:
+            self.float_bias = self.bias.detach().clone()
+
+        if self.weight.dtype == torch.float32:
+            _precision = 2 ** (23 + 1)
+        elif self.weight.dtype == torch.float64:
+            _precision = 2 ** (52 + 1)
+        else:
+            logging.warning(
+                f"Unsupported dtype {self.weight.dtype}. Behaviour may lead unexpected results."
+            )
+            _precision = 2 ** (23 + 1)
+
+        ###### REFERENCE FROM VCMRS ######
+        #sf const
+        sf_const = 48
+
+        #N = np.prod(self.weight.shape[1:])
+        N = np.prod(self.weight.shape) / self.weight.shape[1] # (in, out, kH, kW)
+        self.N = N
+        self.factor = np.sqrt(_precision)
+        #self.sf = 1/6 #precision bits allocation factor
+        self.sf = np.sqrt(sf_const / N)
+
+        # perform the calculate ion CPU to stabalize the calculation
+        self.w_sum = self.weight.cpu().abs().sum(axis=[0, 2, 3]).to(self.weight.device)
+        self.w_sum[self.w_sum == 0] = 1 # prevent divide by 0
+
+        self.fw = (self.factor / self.sf -  np.sqrt(N / 12) * 5) / self.w_sum
+
+        # intify weights
+        self.weight.requires_grad = False  # Just make sure
+        self.weight.copy_(
+            torch.round(self.weight.detach().clone() * self.fw.view(1, -1, 1, 1))
+        )
+
+        # set bias to 0
+        if self.bias is not None:
+            self.bias.requires_grad = False  # Just make sure
+            self.bias.zero_()
+
+        ###### END OF REFERENCE FROM VCMRS ######
+
+    def integer_transposeconv2d(self, x: torch.Tensor):
+        _dtype = x.dtype
+        _cudnn_enabled = torch.backends.cudnn.enabled
+        torch.backends.cudnn.enabled = False
+
+        ###### REFERENCE FROM VCMRS ######
+        # Calculate factor
+        fx = 1
+
+        x_abs = x.abs()
+        x_max = x_abs.max()
+        if x_max > 0:
+            fx = (self.factor * self.sf - 0.5) / x_max
+
+        # intify x
+        x = torch.round(fx * x)
+        x = super().forward(x)
+
+        # x should be all integers
+        x /= fx * self.fw.view(-1, 1, 1)
+        x = x.float()
+
+        # apply bias in float format
+        x = (x.permute(0, 2, 3, 1) + self.float_bias).permute(0, 3, 1, 2).contiguous()
+
+        ###### REFERENCE FROM VCMRS ######
+        torch.backends.cudnn.enabled = _cudnn_enabled
+
+        return out_x.to(_dtype)