Move freeze_at to config; fix an issue of GN from-scratch

Summary:
Move freeze_at to config.

Fix GN from-scratch scripts with freeze_at=0.

Reviewed By: rbgirshick

Differential Revision: D8325882

fbshipit-source-id: 9f362cb269ffd962056dfe8adc888083cff3d9b2

configs/04_2018_gn_baselines/scratch_e2e_mask_rcnn_R-101-FPN_3x_gn.yaml

@@ -37,6 +37,7 @@ MRCNN:
   CONV_INIT: MSRAFill  # default GaussianFill
 TRAIN:
   # WEIGHTS: N/A
+  FREEZE_AT: 0
   DATASETS: ('coco_2014_train', 'coco_2014_valminusminival')
   SCALES: (800,)
   MAX_SIZE: 1333

configs/04_2018_gn_baselines/scratch_e2e_mask_rcnn_R-50-FPN_3x_gn.yaml

@@ -37,6 +37,7 @@ MRCNN:
   CONV_INIT: MSRAFill  # default GaussianFill
 TRAIN:
   # WEIGHTS: N/A
+  FREEZE_AT: 0
   DATASETS: ('coco_2014_train', 'coco_2014_valminusminival')
   SCALES: (800,)
   MAX_SIZE: 1333

detectron/core/config.py

@@ -189,6 +189,9 @@ __C.TRAIN.FREEZE_CONV_BODY = False
 # output directory
 __C.TRAIN.AUTO_RESUME = True
 
+# Add StopGrad at a specified stage so the bottom layers are frozen
+__C.TRAIN.FREEZE_AT = 2
+
 
 # ---------------------------------------------------------------------------- #
 # Data loader options (see detectron/roi_data/loader.py for more info)

detectron/modeling/ResNet.py

@@ -88,10 +88,11 @@ def add_stage(
     return blob_in, dim_in
 
 
-def add_ResNet_convX_body(model, block_counts, freeze_at=2):
+def add_ResNet_convX_body(model, block_counts):
     """Add a ResNet body from input data up through the res5 (aka conv5) stage.
     The final res5/conv5 stage may be optionally excluded (hence convX, where
     X = 4 or 5)."""
+    freeze_at = cfg.TRAIN.FREEZE_AT
     assert freeze_at in [0, 2, 3, 4, 5]
 
     # add the stem (by default, conv1 and pool1 with bn; can support gn)