[feat] kmac calculation for yolox_darknet53

Author: jianwensong

compressai_vision/model_wrappers/base_wrapper.py

@@ -65,6 +65,10 @@ def forward(self, x, input_map_function):
         """Complete the downstream task with end-to-end manner all the way from the input"""
         raise NotImplementedError
 
+    def calc_complexity(self, mode, input, data):
+        """Computes the MACs Complexity of the model"""
+        raise NotImplementedError
+
     @property
     def cfg(self):
         return None

compressai_vision/model_wrappers/detectron2.py

@@ -36,6 +36,10 @@
 
 from compressai_vision.registry import register_vision_model
 
+from ..utils.measure_complexity import (
+    calc_complexity_nn_part1_plyr,
+    calc_complexity_nn_part2_plyr,
+)
 from .base_wrapper import BaseWrapper
 from .intconv2d import IntConv2d, IntTransposedConv2d
 
@@ -559,6 +563,17 @@ def forward(self, x):
     def cfg(self):
         return self._cfg
 
+    def calc_complexity(self, mode, input, data=None):
+        """Computes the MACs Complexity of the model"""
+        if mode == "nn_part_1":
+            return calc_complexity_nn_part1_plyr(self, input)
+        elif mode == "nn_part_2":
+            return calc_complexity_nn_part2_plyr(self, input, data)
+        else:
+            raise NotImplementedError(
+                f"Complexity calculation for {mode} not implemented for Detectron2"
+            )
+
 
 @register_vision_model("faster_rcnn_X_101_32x8d_FPN_3x")
 class faster_rcnn_X_101_32x8d_FPN_3x(Rcnn_R_50_X_101_FPN):

compressai_vision/model_wrappers/jde.py

@@ -36,6 +36,10 @@
 
 from compressai_vision.registry import register_vision_model
 
+from ..utils.measure_complexity import (
+    calc_complexity_nn_part1_dn53,
+    calc_complexity_nn_part2_dn53,
+)
 from .base_wrapper import BaseWrapper
 
 __all__ = [
@@ -489,3 +493,14 @@ def forward(self, x):
                 online_ids.append(tid)
 
         return {"tlwhs": online_tlwhs, "ids": online_ids}
+
+    def calc_complexity(self, mode, input, data=None):
+        """Computes the MACs Complexity of the model"""
+        if mode == "nn_part_1":
+            return calc_complexity_nn_part1_dn53(self, input)
+        elif mode == "nn_part_2":
+            return calc_complexity_nn_part2_dn53(self, input)
+        else:
+            raise NotImplementedError(
+                f"Complexity calculation for {mode} not implemented for JDE"
+            )

compressai_vision/model_wrappers/yolox.py

@@ -36,6 +36,10 @@
 
 from compressai_vision.registry import register_vision_model
 
+from ..utils.measure_complexity import (
+    calc_complexity_nn_part1_yolox,
+    calc_complexity_nn_part2_yolox,
+)
 from .base_wrapper import BaseWrapper
 from .split_squeezes import squeeze_yolox
 
@@ -326,3 +330,14 @@ def forward(self, x):
         )
 
         return pred
+
+    def calc_complexity(self, mode, input, data=None):
+        """Computes the MACs Complexity of the model"""
+        if mode == "nn_part_1":
+            return calc_complexity_nn_part1_yolox(self, input)
+        elif mode == "nn_part_2":
+            return calc_complexity_nn_part2_yolox(self, input)
+        else:
+            raise NotImplementedError(
+                f"Complexity calculation for {mode} not implemented for YOLOX-Darknet53"
+            )

compressai_vision/pipelines/split_inference/image_split_inference.py

@@ -111,7 +111,7 @@ def __call__(
                     break
 
                 if self.is_mac_calculation:
-                    macs, pixels = calc_complexity_nn_part1_plyr(vision_model, d)
+                    macs, pixels = vision_model.calc_complexity("nn_part_1", d)
                     self.acc_kmac_and_pixels_info("nn_part_1", macs, pixels)
 
                 start = time_measure()
@@ -200,8 +200,8 @@ def __call__(
             dec_features["file_name"] = d[0]["file_name"]
             dec_features["file_origin"] = d[0]["file_name"]
             if self.is_mac_calculation:
-                macs, pixels = calc_complexity_nn_part2_plyr(
-                    vision_model, dec_features["data"], dec_features
+                macs, pixels = vision_model.calc_complexity(
+                    "nn_part_2", dec_features, dec_features["data"]
                 )
                 self.acc_kmac_and_pixels_info("nn_part_2", macs, pixels)
 

compressai_vision/pipelines/split_inference/video_split_inference.py

@@ -140,10 +140,7 @@ def __call__(
                     break
 
                 if self.is_mac_calculation and e == self._codec_skip_n_frames:
-                    if hasattr(vision_model, "darknet"):  # for jde
-                        kmacs, pixels = calc_complexity_nn_part1_dn53(vision_model, d)
-                    else:  # for detectron2
-                        kmacs, pixels = calc_complexity_nn_part1_plyr(vision_model, d)
+                    kmacs, pixels = vision_model.calc_complexity("nn_part_1", d)
                     self.add_kmac_and_pixels_info("nn_part_1", kmacs, pixels)
 
                 start = time_measure()
@@ -299,14 +296,9 @@ def __call__(
             )  # Assuming one qp will be used
 
             if self.is_mac_calculation and e == 0:
-                if hasattr(vision_model, "darknet"):  # for jde
-                    kmacs, pixels = calc_complexity_nn_part2_dn53(
-                        vision_model, dec_features
-                    )
-                else:  # for detectron2
-                    kmacs, pixels = calc_complexity_nn_part2_plyr(
-                        vision_model, data, dec_features
-                    )
+                kmacs, pixels = vision_model.calc_complexity(
+                    "nn_part_2", dec_features, data
+                )
                 self.add_kmac_and_pixels_info("nn_part_2", kmacs, pixels)
 
             start = time_measure()

compressai_vision/utils/measure_complexity.py

@@ -3,6 +3,7 @@
 from functools import reduce
 
 import torch
+import torch.nn as nn
 
 from ptflops import get_model_complexity_info
 
@@ -72,7 +73,7 @@ def calc_complexity_nn_part1_plyr(vision_model, img):
     return kmacs, pixels
 
 
-def calc_complexity_nn_part2_plyr(vision_model, data, dec_features):
+def calc_complexity_nn_part2_plyr(vision_model, dec_features, data):
     if isinstance(data[0], list):  # image task
         data = {k: v[0] for k, v in data.items()}
 
@@ -147,6 +148,120 @@ def get_downsampled_shape(h, w, ratio):
     return h, w
 
 
+class YoloxPart1(nn.Module):
+    def __init__(self, vision_model, split_id):
+        super().__init__()
+        self.backbone = vision_model.backbone
+        self.split_id = split_id
+        self.squeeze_at_split_enabled = vision_model.squeeze_at_split_enabled
+        if self.squeeze_at_split_enabled:
+            self.squeeze_model = vision_model.squeeze_model
+
+    def forward(self, x):
+        if self.split_id == "l13":
+            y = self.backbone.stem(x)
+            y = self.backbone.dark2(y)
+            y = self.backbone.dark3[0](y)
+            if self.squeeze_at_split_enabled:
+                y = self.squeeze_model.squeeze_(y)
+        elif self.split_id == "l37":
+            y = self.backbone.stem(x)
+            y = self.backbone.dark2(y)
+            y = self.backbone.dark3(y)
+        return y
+
+
+class YoloxPart2(nn.Module):
+    def __init__(self, vision_model, split_id):
+        super().__init__()
+        self.backbone = vision_model.backbone
+        self.out1_cbl = vision_model.yolo_fpn.out1_cbl
+        self.out1 = vision_model.yolo_fpn.out1
+        self.out2_cbl = vision_model.yolo_fpn.out2_cbl
+        self.out2 = vision_model.yolo_fpn.out2
+        self.upsample = vision_model.yolo_fpn.upsample
+        self.head = vision_model.head
+        self.split_id = split_id
+        self.squeeze_at_split_enabled = vision_model.squeeze_at_split_enabled
+        if self.squeeze_at_split_enabled:
+            self.squeeze_model = vision_model.squeeze_model
+        # self.postprocess = vision_model.postprocess # Not needed for MAC calc
+
+    def forward(self, x):
+        y = x
+        if self.split_id == "l13":
+            if self.squeeze_at_split_enabled:
+                y = self.squeeze_model.expand_(y)
+            for proc_module in self.backbone.dark3[1:]:
+                y = proc_module(y)
+
+        fp_lvl2 = y
+        fp_lvl1 = self.backbone.dark4(fp_lvl2)
+        fp_lvl0 = self.backbone.dark5(fp_lvl1)
+
+        # yolo branch 1
+        b1_in = self.out1_cbl(fp_lvl0)
+        b1_in = self.upsample(b1_in)
+        b1_in = torch.cat([b1_in, fp_lvl1], 1)
+        fp_lvl1 = self.out1(b1_in)
+
+        # yolo branch 2
+        b2_in = self.out2_cbl(fp_lvl1)
+        b2_in = self.upsample(b2_in)
+        b2_in = torch.cat([b2_in, fp_lvl2], 1)
+        fp_lvl2 = self.out2(b2_in)
+
+        outputs = self.head((fp_lvl2, fp_lvl1, fp_lvl0))
+        return outputs
+
+
+def calc_complexity_nn_part1_yolox(vision_model, img):
+    device = torch.device(vision_model.device)
+    img = img[0]["image"].unsqueeze(0).to(device)
+
+    partial_model = YoloxPart1(vision_model, vision_model.split_id)
+
+    C, H, W = img.shape[1:]
+
+    kmacs, _ = measure_mac(
+        partial_model=partial_model,
+        input_res=(C, H, W),
+        input_constructor=None,
+    )
+
+    pixels = reduce(operator.mul, [p_size for p_size in img.shape])
+    return kmacs, pixels
+
+
+def calc_complexity_nn_part2_yolox(vision_model, dec_features):
+    assert "data" in dec_features
+
+    x_data = dec_features["data"]
+
+    x_data = {
+        k: (v[0] if isinstance(x_data[0], list) else v).to(vision_model.device)
+        for k, v in zip(vision_model.split_layer_list, x_data.values())
+    }
+
+    input_tensor = x_data[vision_model.split_id]
+
+    if input_tensor.dim() == 3:
+        input_tensor = input_tensor.unsqueeze(0)
+
+    C, H, W = input_tensor.shape[1:]
+    partial_model = YoloxPart2(vision_model, vision_model.split_id)
+
+    kmacs, _ = measure_mac(
+        partial_model=partial_model,
+        input_res=(C, H, W),
+        input_constructor=None,
+    )
+
+    pixels = reduce(operator.mul, input_tensor.shape)
+
+    return kmacs, pixels
+
+
 def prepare_proposal_input_fpn(resolutions):
     b, c, h, w = resolutions[1]
     resized_img = resolutions[0]