Make SequentialDict return outputs from all modules as DotDict

mart/configs/attack/gain/rcnn_training_loss.yaml

@@ -2,8 +2,8 @@ _target_: mart.nn.CallWith
 module:
   _target_: mart.nn.Sum
 arg_keys:
-  - rpn_loss.loss_objectness
-  - rpn_loss.loss_rpn_box_reg
-  - box_loss.loss_classifier
-  - box_loss.loss_box_reg
+  - "losses_and_detections.training.loss_objectness"
+  - "losses_and_detections.training.loss_rpn_box_reg"
+  - "losses_and_detections.training.loss_classifier"
+  - "losses_and_detections.training.loss_box_reg"
 kwarg_keys: null

mart/configs/model/classifier.yaml

@@ -17,14 +17,6 @@ training_sequence:
   seq040:
     preds:
       _call_with_args_: ["logits"]
-  seq050:
-    output:
-      {
-        "preds": "preds",
-        "target": "target",
-        "logits": "logits",
-        "loss": "loss",
-      }
 
 # The kwargs-centric version.
 # We may use *args as **kwargs to avoid the lengthy _call_with_args_.
@@ -36,10 +28,6 @@ validation_sequence:
   - logits: ["preprocessor"]
   - preds:
       input: logits
-  - output:
-      preds: preds
-      target: target
-      logits: logits
 
 # The simplified version.
 #   We treat a list as the `_call_with_args_` parameter.
@@ -50,8 +38,6 @@ test_sequence:
     logits: ["preprocessor"]
   seq030:
     preds: ["logits"]
-  seq040:
-    output: { preds: preds, target: target, logits: logits }
 
 modules:
   preprocessor: ???
@@ -64,6 +50,3 @@ modules:
   preds:
     _target_: torch.nn.Softmax
     dim: 1
-
-  output:
-    _target_: mart.nn.ReturnKwargs

mart/configs/model/torchvision_faster_rcnn.yaml

@@ -4,13 +4,10 @@ defaults:
 
 # log all losses separately in training.
 training_step_log:
-  [
-    "rpn_loss.loss_objectness",
-    "rpn_loss.loss_rpn_box_reg",
-    "box_loss.loss_classifier",
-    "box_loss.loss_box_reg",
-    "loss",
-  ]
+  loss_objectness: "losses_and_detections.training.loss_objectness"
+  loss_rpn_box_reg: "losses_and_detections.training.loss_rpn_box_reg"
+  loss_classifier: "losses_and_detections.training.loss_classifier"
+  loss_box_reg: "losses_and_detections.training.loss_box_reg"
 
 training_sequence:
   seq010:
@@ -29,55 +26,20 @@ training_sequence:
         "losses_and_detections.training.loss_box_reg",
       ]
 
-  seq040:
-    output:
-      # Output all losses for logging, defined in model.training_step_log
-      {
-        "preds": "losses_and_detections.eval",
-        "target": "target",
-        "loss": "loss",
-        "rpn_loss.loss_objectness": "losses_and_detections.training.loss_objectness",
-        "rpn_loss.loss_rpn_box_reg": "losses_and_detections.training.loss_rpn_box_reg",
-        "box_loss.loss_classifier": "losses_and_detections.training.loss_classifier",
-        "box_loss.loss_box_reg": "losses_and_detections.training.loss_box_reg",
-      }
-
 validation_sequence:
   seq010:
     preprocessor: ["input"]
 
   seq020:
     losses_and_detections: ["preprocessor", "target"]
 
-  seq030:
-    output:
-      {
-        "preds": "losses_and_detections.eval",
-        "target": "target",
-        "rpn_loss.loss_objectness": "losses_and_detections.training.loss_objectness",
-        "rpn_loss.loss_rpn_box_reg": "losses_and_detections.training.loss_rpn_box_reg",
-        "box_loss.loss_classifier": "losses_and_detections.training.loss_classifier",
-        "box_loss.loss_box_reg": "losses_and_detections.training.loss_box_reg",
-      }
-
 test_sequence:
   seq010:
     preprocessor: ["input"]
 
   seq020:
     losses_and_detections: ["preprocessor", "target"]
 
-  seq030:
-    output:
-      {
-        "preds": "losses_and_detections.eval",
-        "target": "target",
-        "rpn_loss.loss_objectness": "losses_and_detections.training.loss_objectness",
-        "rpn_loss.loss_rpn_box_reg": "losses_and_detections.training.loss_rpn_box_reg",
-        "box_loss.loss_classifier": "losses_and_detections.training.loss_classifier",
-        "box_loss.loss_box_reg": "losses_and_detections.training.loss_box_reg",
-      }
-
 modules:
   losses_and_detections:
     # 17s: DualModeGeneralizedRCNN

mart/configs/model/torchvision_object_detection.yaml

@@ -3,14 +3,15 @@ defaults:
   - modular
   - /model/modules@modules.preprocessor: tuple_normalizer
 
-training_step_log: ???
+training_step_log:
+  loss: "loss"
 
 training_sequence: ???
-
 validation_sequence: ???
-
 test_sequence: ???
 
+output_preds_key: "losses_and_detections.eval"
+
 modules:
   losses_and_detections:
     # Return losses in the training mode and predictions in the eval mode in one pass.
@@ -19,6 +20,3 @@ modules:
 
   loss:
     _target_: mart.nn.Sum
-
-  output:
-    _target_: mart.nn.ReturnKwargs

mart/configs/model/torchvision_retinanet.yaml

@@ -3,7 +3,9 @@ defaults:
   - torchvision_object_detection
 
 # log all losses separately in training.
-training_step_log: ["loss_classifier", "loss_box_reg"]
+training_step_log:
+  loss_classifier: "losses_and_detections.training.classification"
+  loss_box_reg: "losses_and_detections.training.bbox_regression"
 
 training_sequence:
   - preprocessor: ["input"]
@@ -14,37 +16,14 @@ training_sequence:
         "losses_and_detections.training.classification",
         "losses_and_detections.training.bbox_regression",
       ]
-  - output:
-      # Output all losses for logging, defined in model.training_step_log
-      {
-        "preds": "losses_and_detections.eval",
-        "target": "target",
-        "loss": "loss",
-        "loss_classifier": "losses_and_detections.training.classification",
-        "loss_box_reg": "losses_and_detections.training.bbox_regression",
-      }
 
 validation_sequence:
   - preprocessor: ["input"]
   - losses_and_detections: ["preprocessor", "target"]
-  - output:
-      {
-        "preds": "losses_and_detections.eval",
-        "target": "target",
-        "loss_classifier": "losses_and_detections.training.classification",
-        "loss_box_reg": "losses_and_detections.training.bbox_regression",
-      }
 
 test_sequence:
   - preprocessor: ["input"]
   - losses_and_detections: ["preprocessor", "target"]
-  - output:
-      {
-        "preds": "losses_and_detections.eval",
-        "target": "target",
-        "loss_classifier": "losses_and_detections.training.classification",
-        "loss_box_reg": "losses_and_detections.training.bbox_regression",
-      }
 
 modules:
   losses_and_detections:

mart/models/modular.py

@@ -134,19 +134,15 @@ def training_step(self, batch, batch_idx):
         for log_name, output_key in self.training_step_log.items():
             self.log(f"training/{log_name}", output[output_key])
 
-        assert "loss" in output
-        return output
-
-    def training_step_end(self, output):
         if self.training_metrics is not None:
             # Some models only return loss in the training mode.
             if self.output_preds_key not in output or self.output_target_key not in output:
                 raise ValueError(
                     f"You have specified training_metrics, but the model does not return {self.output_preds_key} or {self.output_target_key} during training. You can either nullify training_metrics or configure the model to return {self.output_preds_key} and {self.output_target_key} in the training output."
                 )
             self.training_metrics(output[self.output_preds_key], output[self.output_target_key])
-        loss = output.pop(self.output_loss_key)
-        return loss
+
+        return output[self.output_loss_key]
 
     def training_epoch_end(self, outputs):
         if self.training_metrics is not None:
@@ -168,13 +164,9 @@ def validation_step(self, batch, batch_idx):
         for log_name, output_key in self.validation_step_log.items():
             self.log(f"validation/{log_name}", output[output_key])
 
-        return output
-
-    def validation_step_end(self, output):
         self.validation_metrics(output[self.output_preds_key], output[self.output_target_key])
 
-        # I don't know why this is required to prevent CUDA memory leak in validaiton and test. (Not required in training.)
-        output.clear()
+        return None
 
     def validation_epoch_end(self, outputs):
         metrics = self.validation_metrics.compute()
@@ -194,13 +186,9 @@ def test_step(self, batch, batch_idx):
         for log_name, output_key in self.test_step_log.items():
             self.log(f"test/{log_name}", output[output_key])
 
-        return output
-
-    def test_step_end(self, output):
         self.test_metrics(output[self.output_preds_key], output[self.output_target_key])
 
-        # I don't know why this is required to prevent CUDA memory leak in validaiton and test. (Not required in training.)
-        output.clear()
+        return None
 
     def test_epoch_end(self, outputs):
         metrics = self.test_metrics.compute()

mart/nn/nn.py

@@ -49,10 +49,6 @@ class SequentialDict(torch.nn.ModuleDict):
     """
 
     def __init__(self, modules, sequences=None):
-
-        if "output" not in modules:
-            raise ValueError("Modules must have an module named 'output'")
-
         super().__init__(modules)
 
         self._sequences = {
@@ -121,7 +117,8 @@ def forward(self, step=None, sequence=None, **kwargs):
             # Pop the executed module to proceed with the sequence
             sequence.popitem(last=False)
 
-        return kwargs["output"]
+        # return kwargs as DotDict
+        return DotDict(kwargs)
 
 
 class ReturnKwargs(torch.nn.Module):

tests/test_experiments.py

@@ -209,7 +209,7 @@ def test_coco_fasterrcnn_experiment(coco_cfg, tmp_path):
         "-m",
         "experiment=COCO_TorchvisionFasterRCNN",
         "hydra.sweep.dir=" + str(tmp_path),
-        "optimized_metric=training/rpn_loss.loss_objectness",
+        "optimized_metric=training/loss_objectness",
     ] + overrides
     run_sh_command(command)
 
@@ -224,7 +224,7 @@ def test_coco_fasterrcnn_adv_experiment(coco_cfg, tmp_path):
         "-m",
         "experiment=COCO_TorchvisionFasterRCNN_Adv",
         "hydra.sweep.dir=" + str(tmp_path),
-        "optimized_metric=training/rpn_loss.loss_objectness",
+        "optimized_metric=training/loss_objectness",
     ] + overrides
     run_sh_command(command)
 
@@ -256,7 +256,7 @@ def test_armory_carla_fasterrcnn_experiment(carla_cfg, tmp_path):
         "experiment=ArmoryCarlaOverObjDet_TorchvisionFasterRCNN",
         "+attack@model.modules.input_adv_test=object_detection_mask_adversary",
         "hydra.sweep.dir=" + str(tmp_path),
-        "optimized_metric=training/rpn_loss.loss_objectness",
+        "optimized_metric=training/loss_objectness",
     ] + overrides
     run_sh_command(command)