Make MMDetection config object accessible to users

docs/mkdocs.yml

@@ -58,6 +58,7 @@ nav:
     - How to use negative samples: negative_samples.md
     - Fixed Splitter: voc_predefined_splits.md
     - Progressive resizing: progressive_resizing.md
+    - MMDetection Custom Config: mmdet_custom_config.md
   - Transforms:
     - Albumentations: albumentations.md 
   - Models: models.md

icevision/models/mmdet/common/bbox/single_stage/model.py

@@ -14,6 +14,7 @@ def model(
     num_classes: int,
     checkpoints_path: Optional[Union[str, Path]] = "checkpoints",
     force_download=False,
+    cfg_options=None,
 ) -> nn.Module:
 
     return build_model(
@@ -23,4 +24,5 @@ def model(
         pretrained=backbone.pretrained,
         checkpoints_path=checkpoints_path,
         force_download=force_download,
+        cfg_options=cfg_options,
     )

icevision/models/mmdet/common/utils.py

@@ -36,13 +36,15 @@ def build_model(
     pretrained: bool = True,
     checkpoints_path: Optional[Union[str, Path]] = "checkpoints",
     force_download=False,
+    cfg_options=None,
 ) -> nn.Module:
 
     cfg, weights_path = create_model_config(
         backbone=backbone,
         pretrained=pretrained,
         checkpoints_path=checkpoints_path,
         force_download=force_download,
+        cfg_options=cfg_options,
     )
 
     if model_type == "one_stage_detector_bbox":
@@ -70,5 +72,7 @@ def build_model(
         load_checkpoint(_model, str(weights_path))
 
     _model.param_groups = MethodType(param_groups, _model)
+    _model.cfg = cfg  # save the config in the model for convenience
+    _model.weights_path = weights_path  # save the model.weights_path in case we want to rebuild the model after updating its attributes
 
     return _model

icevision/models/mmdet/utils.py

@@ -61,6 +61,7 @@ def create_model_config(
     pretrained: bool = True,
     checkpoints_path: Optional[Union[str, Path]] = "checkpoints",
     force_download=False,
+    cfg_options=None,
 ):
 
     model_name = backbone.model_name
@@ -81,4 +82,7 @@ def create_model_config(
 
     cfg = Config.fromfile(config_path)
 
+    if cfg_options is not None:
+        cfg.merge_from_dict(cfg_options)
+
     return cfg, weights_path

tests/models/mmdet/test_model.py

@@ -3,14 +3,25 @@
 
 
 @pytest.mark.parametrize(
-    "ds, model_type, pretrained",
+    "ds, model_type, pretrained, cfg_options",
     (
-        ("fridge_ds", models.mmdet.retinanet, True),
-        ("fridge_ds", models.mmdet.retinanet, False),
+        ("fridge_ds", models.mmdet.retinanet, True, None),
+        ("fridge_ds", models.mmdet.retinanet, False, None),
+        (
+            "fridge_ds",
+            models.mmdet.retinanet,
+            False,
+            {
+                "model.bbox_head.loss_bbox.loss_weight": 2,
+                "model.bbox_head.loss_cls.loss_weight": 0.8,
+            },
+        ),
     ),
 )
 class TestBboxModels:
-    def dls_model(self, ds, model_type, pretrained, samples_source, request):
+    def dls_model(
+        self, ds, model_type, pretrained, cfg_options, samples_source, request
+    ):
         train_ds, valid_ds = request.getfixturevalue(ds)
         train_dl = model_type.train_dl(train_ds, batch_size=2)
         valid_dl = model_type.valid_dl(valid_ds, batch_size=2)
@@ -19,16 +30,18 @@ def dls_model(self, ds, model_type, pretrained, samples_source, request):
         backbone.config_path = samples_source / backbone.config_path
 
         model = model_type.model(
-            backbone=backbone(pretrained=pretrained), num_classes=5
+            backbone=backbone(pretrained=pretrained),
+            num_classes=5,
+            cfg_options=cfg_options,
         )
 
         return train_dl, valid_dl, model
 
     def test_mmdet_bbox_models_fastai(
-        self, ds, model_type, pretrained, samples_source, request
+        self, ds, model_type, pretrained, cfg_options, samples_source, request
     ):
         train_dl, valid_dl, model = self.dls_model(
-            ds, model_type, pretrained, samples_source, request
+            ds, model_type, pretrained, cfg_options, samples_source, request
         )
 
         learn = model_type.fastai.learner(
@@ -37,10 +50,10 @@ def test_mmdet_bbox_models_fastai(
         learn.fine_tune(1, 3e-4)
 
     def test_mmdet_bbox_models_light(
-        self, ds, model_type, pretrained, samples_source, request
+        self, ds, model_type, pretrained, cfg_options, samples_source, request
     ):
         train_dl, valid_dl, model = self.dls_model(
-            ds, model_type, pretrained, samples_source, request
+            ds, model_type, pretrained, cfg_options, samples_source, request
         )
 
         class LitModel(model_type.lightning.ModelAdapter):