From b6a083237ef3d2afdeb80def681ab942fe6a3117 Mon Sep 17 00:00:00 2001
From: Tuan Tran <1254753+antoine-tran@users.noreply.github.com>
Date: Thu, 13 Jul 2023 13:35:52 -0700
Subject: [PATCH] Enable operators fusion (Conv + BN + ReLu) in FBNet V2

Summary:
X-link: https://github.com/facebookresearch/d2go/pull/592

Pull Request resolved: https://github.com/facebookresearch/mobile-vision/pull/172

The FBNet V2 consists of several basic blocks of pattern conv(2d) + BatchNorm + ReLU. These operators can be fused into one operator, greatly speeding up the computation in GPU.

This diff / PR adds option to build an FBNet V2 backbone in fused manner, i.e. it looks up recursively the blocks and sub-blocks for the known patterns, and fuse them using `mobile_cv.arch.utils.fuse_utils`

Differential Revision: D47395999

fbshipit-source-id: 0f52567d7448ef30151422937a8af5efd6c44645
---
 mobile_cv/arch/builder/meta_builder.py        | 22 ++++++++++++++++++-
 mobile_cv/arch/fbnet_v2/blocks_factory.py     |  1 -
 .../arch/tests/test_builder_meta_builder.py   |  9 ++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/mobile_cv/arch/builder/meta_builder.py b/mobile_cv/arch/builder/meta_builder.py
index b3ecc425..e1f7f57b 100755
--- a/mobile_cv/arch/builder/meta_builder.py
+++ b/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
 
diff --git a/mobile_cv/arch/fbnet_v2/blocks_factory.py b/mobile_cv/arch/fbnet_v2/blocks_factory.py
index 3398bae5..240f7310 100644
--- a/mobile_cv/arch/fbnet_v2/blocks_factory.py
+++ b/mobile_cv/arch/fbnet_v2/blocks_factory.py
@@ -13,7 +13,6 @@
     irf_block,
     mobileone_block,
     res_block,
-    sg_block,
 )
 from torch import nn
 
diff --git a/mobile_cv/arch/tests/test_builder_meta_builder.py b/mobile_cv/arch/tests/test_builder_meta_builder.py
index 85405617..1a01d319 100755
--- a/mobile_cv/arch/tests/test_builder_meta_builder.py
+++ b/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}