Expand the search space; smoothing window for tensorboard metrics

Summary:
Pull Request resolved: #187

Expand the search space to cover kernel size, expansion factor, blocks per stage; smoothing window for tensorboard metrics

Reviewed By: ajinkya-deogade

Differential Revision: D49247666

fbshipit-source-id: 3b796c21e0fb6ba4c78f57cf2aa2298e8ad32999

mobile_cv/model_zoo/models/fbnet_v2.py

@@ -154,17 +154,58 @@ def apply_overwrite_options(
             return arch_def
 
         arch_def = copy.deepcopy(arch_def)
+
+        # Check if we are modulating the number of blocks
+        num_blocks_overwrites = {}
+        for overwrite in overwrite_options:
+            stage, block, value, overwrite_type = (
+                overwrite["STAGE"],
+                overwrite["BLOCK"],
+                overwrite["VALUE"],
+                overwrite["TYPE"],
+            )
+
+            if overwrite_type == "NUM_BLOCKS_IN_STAGE":
+                assert (
+                    stage not in num_blocks_overwrites
+                ), f"multiple number of blocks overwrites for stage {stage}"
+
+                num_blocks_overwrites[stage] = value
+
+        for stage in num_blocks_overwrites.keys():
+            num_blocks = num_blocks_overwrites[stage]
+
+            for ind, entry in enumerate(arch_def["blocks"]):
+                if arch_def["blocks"][ind] is None:
+                    continue
+
+                if entry["stage_idx"] == stage and entry["block_idx"] >= num_blocks:
+                    arch_def["blocks"][ind] = None
+
+        # Delete all elements from arch_def which are None
+        while True:
+            if any(a is None for a in arch_def["blocks"]):
+                arch_def["blocks"].remove(None)
+            else:
+                break
+
         for overwrite in overwrite_options:
-            stage, block, num_channels = (
+            stage, block, value, overwrite_type = (
                 overwrite["STAGE"],
                 overwrite["BLOCK"],
                 overwrite["VALUE"],
+                overwrite["TYPE"],
             )
 
             # Find the block with the proper stage / block index
             for entry in arch_def["blocks"]:
                 if entry["stage_idx"] == stage and entry["block_idx"] == block:
-                    entry["block_cfg"]["out_channels"] = num_channels
+                    if overwrite_type == "OUT_CHANNELS":
+                        entry["block_cfg"]["out_channels"] = value
+                    elif overwrite_type == "KERNEL_SIZE":
+                        entry["block_cfg"]["kernel_size"] = value
+                    elif overwrite_type == "EXPANSION":
+                        entry["block_cfg"]["expansion"] = value
         return arch_def
 
     def forward(self, x):