Fuse fbnet_v2 ConvBNReLu

mobile_cv/arch/builder/meta_builder.py

@@ -141,6 +141,7 @@
 import mobile_cv.common.misc.iter_utils as iu
 import torch.nn as nn
 from mobile_cv.arch.fbnet_v2.blocks_factory import PRIMITIVES
+from mobile_cv.arch.utils import fuse_utils
 
 
 logger = logging.getLogger(__name__)
@@ -476,6 +477,8 @@ def build_blocks(
         dim_in: Optional[int] = None,
         prefix_name: str = "xif",
         override_missing: Optional[Dict[Any, Any]] = None,
+        fuse_ops: bool = False,
+        is_qat: bool = False,
         **kwargs,
     ):
         """
@@ -486,6 +489,9 @@ def build_blocks(
         override_missing: arguments that override the config in the blocks
                 if the argument in the block is None, otherwise it will be
                 ignored
+        fuse_ops: Whether to fuse the ops of known patterns (e.g. conv + bn + relu)
+                to speed up computation
+        is_qat: whether this is quantization aware training
         """
         assert isinstance(blocks, list) and all(
             isinstance(x, dict) for x in blocks
@@ -508,7 +514,12 @@ def build_blocks(
             block_cfg = block["block_cfg"]
             cur_kwargs = update_with_block_kwargs(copy.deepcopy(kwargs), block)
             nnblock = self.build_block(
-                block_op, block_cfg, override_missing=override_missing, **cur_kwargs
+                block_op,
+                block_cfg,
+                override_missing=override_missing,
+                fuse_ops=fuse_ops,
+                is_qat=is_qat,
+                **cur_kwargs,
             )
             nn_name = f"{prefix_name}{stage_idx}_{block_idx}"
             assert nn_name not in modules, f"{nn_name} existed in {modules}"
@@ -523,6 +534,8 @@ def build_block(
         block_cfg: Dict[str, Any],
         override_missing: Optional[Dict[str, Any]] = None,
         replace_strs: Optional[Dict[str, Any]] = None,
+        fuse_ops: bool = False,
+        is_qat: bool = False,
         **kwargs,
     ):
         """
@@ -534,6 +547,9 @@ def build_block(
             if the values are in '{name}` format
             Example: block_cfg={"out_channels": "{feature_dim}"} and replace_strs={"feature_dim": 20}
             will produce block_cfg={"out_channels": 20}
+        fuse_ops: Whether to fuse the ops of known patterns (e.g. conv + bn + relu)
+            to speed up computation
+        is_qat: whether this is quantization aware training
         """
 
         assert "out_channels" in block_cfg
@@ -572,6 +588,10 @@ def build_block(
             out_channels = self._get_divisible_width(out_channels * width_ratio)
 
         ret = PRIMITIVES.get(block_op)(in_channels, out_channels, **new_kwargs)
+        if fuse_ops:
+            # In training mode, only qat is supported for ops fusion
+            is_qat = ret.training or is_qat
+            ret = fuse_utils.fuse_model(ret, is_qat=is_qat, inplace=True)
         self.last_depth = getattr(ret, "out_channels", out_channels)
         return ret
 

mobile_cv/arch/fbnet_v2/blocks_factory.py

@@ -13,7 +13,6 @@
     irf_block,
     mobileone_block,
     res_block,
-    sg_block,
 )
 from torch import nn
 

mobile_cv/arch/tests/test_builder_meta_builder.py

@@ -53,6 +53,15 @@ def test_fbnet_builder_check_output(self):
         output = model(input)
         self.assertEqual(output.shape, torch.Size([2, 8, 2, 2]))
 
+    def test_fbnet_builder_fuse_ops(self):
+        import torch.ao.nn.intrinsic as nni
+
+        arch_def = {"blocks": [[("conv_k3", 4, 2, 1)]]}
+        model, builder = _build_model(arch_def, fuse_ops=True, dim_in=3)
+        self.assertIsInstance(model[0].conv, nni.ConvBnReLU2d)
+        self.assertIsInstance(model[0].bn, torch.nn.Identity)
+        self.assertIsInstance(model[0].relu, torch.nn.Identity)
+
     def test_fbnet_builder_width_divisor(self):
         e6 = {"expansion": 6}