[feat] support intconv2d for Faster RCNN & Mask RCNN

Author: chyomin06

cfgs/vision_model/default.yaml

@@ -6,24 +6,28 @@ faster_rcnn_R_50_FPN_3x:
   model_path_prefix: ${..model_root_path}
   cfg: "models/detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
   weights: "weights/detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
+  integer_conv_weight: False
   splits : "r2" #, "c2" or "fpn"
 
 faster_rcnn_X_101_32x8d_FPN_3x:
   model_path_prefix: ${..model_root_path}
   cfg: "models/detectron2/configs/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"
   weights: "weights/detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
+  integer_conv_weight: False
   splits : "fpn" #, "c2" or "r2"
 
 mask_rcnn_R_50_FPN_3x:
   model_path_prefix: ${..model_root_path}
   cfg: "models/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
   weights: "weights/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
+  integer_conv_weight: False
   splits : "r2" #, "c2" or "fpn"
 
 mask_rcnn_X_101_32x8d_FPN_3x:
   model_path_prefix: ${..model_root_path}
   cfg: "models/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml"
   weights: "weights/detectron2/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x/139653917/model_final_2d9806.pkl"
+  integer_conv_weight: False
   splits : "fpn" #, "c2" or "r2"
 
 panoptic_rcnn_R_101_FPN_3x:

compressai_vision/model_wrappers/detectron2.py

@@ -27,6 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import re
 from enum import Enum
 from pathlib import Path
 from typing import Dict, List
@@ -45,6 +46,7 @@
 from compressai_vision.registry import register_vision_model
 
 from .base_wrapper import BaseWrapper
+from .intconv2d import IntConv2d
 
 __all__ = [
     "faster_rcnn_X_101_32x8d_FPN_3x",
@@ -58,6 +60,51 @@
 root_path = thisdir.joinpath("../..")
 
 
+# reference Conv2D from detectron2/detercon2/layers/wrappers.py
+class Conv2d(IntConv2d):
+    def __init__(self, *args, **kwargs) -> None:
+        """
+        Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
+
+        Args:
+            norm (nn.Module, optional): a normalization layer
+            activation (callable(Tensor) -> Tensor): a callable activation function
+
+        It assumes that norm layer is used before activation.
+        """
+
+        norm = kwargs.pop("norm", None)
+        activation = kwargs.pop("activation", None)
+        super().__init__(*args, **kwargs)
+
+        self.norm = norm
+        self.activation = activation
+
+    def set_attributes(self, module):
+
+        if hasattr(module, "norm"):
+            self.norm = module.norm
+
+        if hasattr(module, "activation"):
+            self.activation = module.activation
+
+        if hasattr(module, "bias"):
+            self.bias = module.bias
+
+    def forward(self, x: torch.Tensor):
+        if not self.initified_weight_mode:
+            x = self.conv2d(x)
+        else:
+            x = self.integer_conv2d(x)
+
+        if self.norm is not None:
+            x = self.norm(x)
+        if self.activation is not None:
+            x = self.activation(x)
+
+        return x
+
+
 class Split_Points(Enum):
     def __str__(self):
         return str(self.value)
@@ -79,16 +126,28 @@ def __init__(self, device: str, **kwargs):
             else kwargs["model_path_prefix"]
         )
         self._cfg.merge_from_file(f"{_path_prefix}/{kwargs['cfg']}")
+        _integer_conv_weight = bool(kwargs["integer_conv_weight"])
 
-        self.model = build_model(self._cfg).to(device).eval()
+        self.model = build_model(self._cfg)
+        self.replace_conv2d_modules(self.model)
+        self.model = self.model.to(device).eval()
+
+        DetectionCheckpointer(self.model).load(f"{_path_prefix}/{kwargs['weights']}")
+
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        # must be called after loading weights to a model
+        if _integer_conv_weight:
+            self.mode = self.quantize_weights(self.model)
 
         self.backbone = self.model.backbone
         self.top_block = self.model.backbone.top_block
         self.proposal_generator = self.model.proposal_generator
         self.roi_heads = self.model.roi_heads
         self.postprocess = self.model._postprocess
-        DetectionCheckpointer(self.model).load(f"{_path_prefix}/{kwargs['weights']}")
 
+        # to be used for printing info logs
         self.model_info = {"cfg": kwargs["cfg"], "weights": kwargs["weights"]}
 
         self.supported_split_points = Split_Points
@@ -128,6 +187,38 @@ def SPLIT_R2(self):
     def size_divisibility(self):
         return self.backbone.size_divisibility
 
+    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)
+
+                # Since regular list is used instead of ModuleList
+                if "fpn_lateral" in child_name or "fpn_output" in child_name:
+                    idx = re.findall(r"\d", child_name)
+                    assert len(idx) == 1
+                    idx = int(idx[0])
+                    assert idx in [2, 3, 4, 5]
+
+                    if "fpn_lateral" in child_name:
+                        module.lateral_convs[3 - (idx - 2)] = int_conv2d
+                    else:
+                        assert "fpn_output" in child_name
+                        module.output_convs[3 - (idx - 2)] = int_conv2d
+
+                setattr(module, child_name, int_conv2d)
+            else:
+                self.replace_conv2d_modules(child_module)
+
+    @staticmethod
+    def quantize_weights(model):
+        for _, m in model.named_modules():
+            if type(m).__name__ == "Conv2d":
+                # print(f"Module name: {name} and type {type(m).__name__}")
+                m.quantize_weights()
+
+        return model
+
     def input_resize(self, images: List):
         return ImageList.from_tensors(images, self.size_divisibility)
 

…vision/model_wrappers/intconv_wrapper.py …essai_vision/model_wrappers/intconv2d.pycompressai_vision/model_wrappers/intconv_wrapper.py renamed to compressai_vision/model_wrappers/intconv2d.py compressai_vision/model_wrappers/intconv_wrapper.py renamed to compressai_vision/model_wrappers/intconv2d.py

@@ -27,57 +27,27 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import copy
 import logging
 
 import numpy as np
 import torch
-from torch import nn
-
-
-class IntConv2dWrapper(nn.Conv2d):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups: int = 1,
-        bias: bool = True,
-        padding_mode: str = "zeros",
-        device=None,
-        dtype=None,
-    ) -> None:
-        super().__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            dilation,
-            groups,
-            bias,
-            padding_mode,
-            device,
-            dtype,
-        )
-        self.initified_weight_mode = False
+from torch.nn import functional as F
 
-    """
-    def _set_mode(mode):
-        global _precision, _high_precision, _mode
 
-        if mode == 'none':
-            _precision = 0
-        elif mode == 'float32':
-            _precision = 2**(23+1)
-        elif mode == 'float64':
-            _precision = 2**(52+1)
+class IntConv2d(torch.nn.Conv2d):
+    def __init__(self, *args, **kwargs) -> None:
+        _nkwargs = copy.deepcopy(kwargs)
 
-        _mode = mode
-        torch.backends.cudnn.enabled = mode=='none'
-    """
+        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
 
     def quantize_weights(self):
         self.initified_weight_mode = True
@@ -97,7 +67,7 @@ def quantize_weights(self):
             )
             _precision = 2 ** (23 + 1)
 
-        ###### REFERENCE FROM VCMRMS ######
+        ###### REFERENCE FROM VCMRS ######
         # sf const
         sf_const = 48
 
@@ -124,18 +94,14 @@ def quantize_weights(self):
             self.bias.requires_grad = False  # Just make sure
             self.bias.zero_()
 
-        ###### END OF REFERENCE FROM VCMRMS ######
-
-    def forward(self, x: torch.Tensor):
-        if not self.initified_weight_mode:
-            return super().forward(x)
+        ###### END OF REFERENCE FROM VCMRS ######
 
+    def integer_conv2d(self, x: torch.Tensor):
         _dtype = x.dtype
         _cudnn_enabled = torch.backends.cudnn.enabled
         torch.backends.cudnn.enabled = False
 
-        ###### REFERENCE FROM VCMRMS ######
-
+        ###### REFERENCE FROM VCMRS ######
         # Calculate factor
         fx = 1
 
@@ -145,17 +111,39 @@ def forward(self, x: torch.Tensor):
             fx = (self.factor * self.sf - 0.5) / x_max
 
         # intify x
-        x = torch.round(fx * x)
-        x = super().forward(x)
+        out_x = torch.round(fx * x)
+
+        out_x = F.conv2d(
+            out_x,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
 
         # x should be all integers
-        x /= fx * self.fw.view(-1, 1, 1)
-        x = x.float()
+        out_x = out_x / (fx * self.fw.view(-1, 1, 1)).float()
 
         # apply bias in float format
-        x = (x.permute(0, 2, 3, 1) + self.float_bias).permute(0, 3, 1, 2).contiguous()
-        ###### REFERENCE FROM VCMRMS ######
-
+        out_x = (
+            (out_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 x.to(_dtype)
+        return out_x.to(_dtype)
+
+    def conv2d(self, x: torch.Tensor):
+        return F.conv2d(
+            x,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )

compressai_vision/model_wrappers/jde.py

@@ -88,7 +88,7 @@ def __init__(self, device: str, **kwargs):
             self.model_configs["frame_rate"] / 30.0 * self.model_configs["track_buffer"]
         )
 
-        integer_conv_weight = bool(kwargs["integer_conv_weight"])
+        _integer_conv_weight = bool(kwargs["integer_conv_weight"])
 
         assert "splits" in kwargs, "Split layer ids must be provided"
         self.split_layer_list = kwargs["splits"]
@@ -106,7 +106,7 @@ def __init__(self, device: str, **kwargs):
             param.requires_grad = False
 
         # must be called after loading weights to a model
-        if integer_conv_weight:
+        if _integer_conv_weight:
             self.darknet = self.quantize_weights(self.darknet)
 
         self.kalman_filter = KalmanFilter()
@@ -128,10 +128,10 @@ def reset(self):
 
     @staticmethod
     def quantize_weights(model):
-
         for module_def, module in zip(model.module_defs, model.module_list):
             if module_def["type"] == "convolutional":
                 for m in module:
+                    # print(type(m).__name__)
                     if type(m).__name__ == "IntConv2dWrapper":
                         m.quantize_weights()
 

compressai_vision/model_wrappers/jde_lowlevel.py

@@ -27,10 +27,11 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import torch
 import torch.nn as nn
 from jde.models import EmptyLayer, Upsample, YOLOLayer
 
-from .intconv_wrapper import IntConv2dWrapper
+from .intconv2d import IntConv2d
 
 try:
     from jde.utils.syncbn import SyncBN
@@ -40,6 +41,17 @@
     batch_norm = nn.BatchNorm2d
 
 
+class IntConv2dWrapper(IntConv2d):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor):
+        if not self.initified_weight_mode:
+            return self.conv2d(x)
+
+        return self.integer_conv2d(x)
+
+
 def create_modules(module_defs, device: str):
     """
     Constructs module list of layer blocks from module configuration in module_defs