Feat/metrics logging

.gitignore

@@ -28,6 +28,9 @@ __pycache__/
 # vscode sftp settings
 .vscode/sftp.json
 
+# vscode launch settings
+.vscode/launch.json
+
 # redis
 *.rdb
 

.vscode/settings.json

@@ -18,11 +18,6 @@
     // Python use black as formatter
     "editor.defaultFormatter": null
   },
-  "python.formatting.provider": "black",
-  "python.linting.pylintEnabled": true,
-  "python.linting.mypyEnabled": true,
-  "python.linting.pydocstyleEnabled": true,
-  "python.linting.pydocstyleArgs": ["--convention=google"],
   "spellright.language": ["en"],
   "spellright.documentTypes": ["markdown", "latex", "plaintext", "typescript"],
   "files.exclude": {

docs/source/user_guide/faster_rcnn_example.py

@@ -4,6 +4,7 @@
 
 import lightning.pytorch as pl
 import numpy as np
+import tempfile
 from torch.optim import SGD
 from torch.optim.lr_scheduler import LinearLR, MultiStepLR
 
@@ -52,7 +53,11 @@ def get_config() -> ExperimentConfig:
     ######################################################
     ##                    General Config                ##
     ######################################################
-    config = get_default_cfg(exp_name="faster_rcnn_r50_fpn_coco")
+    tmpdir = tempfile.mkdtemp()
+
+    config = get_default_cfg(
+        exp_name="faster_rcnn_r50_fpn_coco", work_dir=tmpdir
+    )
 
     # High level hyper parameters
     params = ExperimentParameters()

requirements/dev.txt

@@ -8,6 +8,8 @@ pydocstyle
 pylint
 pytest
 pytest-notebook
+jupyter
+nbdime==3.2.1
 types-PyYAML
 types-requests
 types-tabulate

tests/engine/loss_module_test.py

@@ -35,16 +35,23 @@ def test_forward(self) -> None:
                         }
                     ),
                 },
-            ]
+            ],
+            exclude_attributes=["dummy_attribute"],
         )
         x = torch.rand(2, 3, 4, 5)
         y = torch.rand(2, 3, 4, 5)
-        losses = loss({"input": x}, {"target": y})
-        total_loss = sum(losses.values())
+        total_loss, metrics = loss({"input": x}, {"target": y})
 
         self.assertAlmostEqual(
             total_loss.item(),
             0.7 * torch.nn.MSELoss()(x, y).item()
             + 0.3 * torch.nn.L1Loss()(x, y).item(),
             places=3,
         )
+
+        self.assertAlmostEqual(
+            metrics["loss"],
+            0.7 * torch.nn.MSELoss()(x, y).detach().cpu().item()
+            + 0.3 * torch.nn.L1Loss()(x, y).detach().cpu().item(),
+            places=3,
+        )

tests/engine/util_test.py

@@ -47,65 +47,69 @@ def test_apply_to_collection():
 
     # test with data as namedtuple or dataclass
     data_cls = Test(1, 2)
-    data = apply_to_collection(data_cls, int, lambda x: x * 2)
-    assert data == Test(2, 4)
+    data_cls = apply_to_collection(data_cls, int, lambda x: x * 2)
+    assert data_cls == Test(2, 4)
 
     data_cls = Test(1, 2)
-    data = apply_to_collection(data_cls, (int, str), lambda x: x * 2)
-    assert data == Test(2, 4)
+    data_cls = apply_to_collection(data_cls, (int, str), lambda x: x * 2)
+    assert data_cls == Test(2, 4)
 
     data_cls = Test(1, 2)
-    data = apply_to_collection(data_cls, int, lambda x: x * 2, wrong_dtype=str)
-    assert data == Test(2, 4)
+    data_cls = apply_to_collection(
+        data_cls, int, lambda x: x * 2, wrong_dtype=str
+    )
+    assert data_cls == Test(2, 4)
 
     data_cls = Test(1, 2)
-    data = apply_to_collection(
+    data_cls = apply_to_collection(
         data_cls,
         int,
         lambda x: x * 2,
         wrong_dtype=(str, int),
         include_none=False,
     )
-    assert data == Test(1, 2)
+    assert data_cls == Test(1, 2)
 
     data_cls = Test(1, 2)
-    data = apply_to_collection(
+    data_cls = apply_to_collection(
         data_cls,
         int,
         lambda x: x * 2,
         wrong_dtype=(str, int),
         include_none=True,
     )
-    assert data == Test(1, 2)
+    assert data_cls == Test(1, 2)
 
     data_tup = namedtuple("test", "aaa bbb")(1, 2)
-    data = apply_to_collection(data_tup, int, lambda x: x * 2)
-    assert data == namedtuple("test", "aaa bbb")(2, 4)
+    data_tup = apply_to_collection(data_tup, int, lambda x: x * 2)
+    assert data_tup == namedtuple("test", "aaa bbb")(2, 4)
 
     data_tup = namedtuple("test", "aaa bbb")(1, 2)
-    data = apply_to_collection(data_tup, (int, str), lambda x: x * 2)
-    assert data == namedtuple("test", "aaa bbb")(2, 4)
+    data_tup = apply_to_collection(data_tup, (int, str), lambda x: x * 2)
+    assert data_tup == namedtuple("test", "aaa bbb")(2, 4)
 
     data_tup = namedtuple("test", "aaa bbb")(1, 2)
-    data = apply_to_collection(data_tup, int, lambda x: x * 2, wrong_dtype=str)
-    assert data == namedtuple("test", "aaa bbb")(2, 4)
+    data_tup = apply_to_collection(
+        data_tup, int, lambda x: x * 2, wrong_dtype=str
+    )
+    assert data_tup == namedtuple("test", "aaa bbb")(2, 4)
 
     data_tup = namedtuple("test", "aaa bbb")(1, 2)
-    data = apply_to_collection(
+    data_tup = apply_to_collection(
         data_tup,
         int,
         lambda x: x * 2,
         wrong_dtype=(str, int),
         include_none=False,
     )
-    assert data == namedtuple("test", "aaa bbb")(1, 2)
+    assert data_tup == namedtuple("test", "aaa bbb")(1, 2)
 
     data_tup = namedtuple("test", "aaa bbb")(1, 2)
-    data = apply_to_collection(
+    data_tup = apply_to_collection(
         data_tup,
         int,
         lambda x: x * 2,
         wrong_dtype=(str, int),
         include_none=True,
     )
-    assert data == namedtuple("test", "aaa bbb")(1, 2)
+    assert data_tup == namedtuple("test", "aaa bbb")(1, 2)

tests/model/detect/mask_rcnn_test.py

@@ -143,15 +143,13 @@ def test_train(self):
                 outputs = mask_rcnn(inputs, images_hw, gt_boxes, gt_class_ids)
                 assert isinstance(outputs, MaskRCNNOut)
 
-                mask_losses = mask_rcnn_loss(outputs, data)
+                total_loss, metrics = mask_rcnn_loss(outputs, data)
 
-                total_loss = sum(mask_losses.values())
                 total_loss.backward()
                 optimizer.step()
 
                 # print statistics
-                losses = {"loss": total_loss, **mask_losses}
-                for k, loss in losses.items():
+                for k, loss in metrics.items():
                     if k in running_losses:
                         running_losses[k] += loss
                     else:

tests/util.py

@@ -242,7 +242,7 @@ def generate_instance_masks(
     torch.random.set_rng_state(torch.manual_seed(0).get_state())
     rand_mask = torch.randint(0, num_masks, (height, width))
     mask_tensor = torch.stack(
-        [(rand_mask == i).type(torch.uint8) for i in range(num_masks)]
+        [(torch.eq(rand_mask, i)).type(torch.uint8) for i in range(num_masks)]
     )
     torch.random.set_rng_state(state)
     return (

vis4d/common/ckpt.py

@@ -325,7 +325,7 @@ def load(module: nn.Module, prefix: str = "") -> None:
         # recursively check parallel module in case that the model has a
         # complicated structure, e.g., nn.Module(nn.Module(DDP))
         if is_module_wrapper(module):
-            module = module.module  # type: ignore
+            module = module.module
         local_metadata = (
             {} if metadata is None else metadata.get(prefix[:-1], {})
         )

vis4d/common/imports.py

@@ -55,8 +55,6 @@ def is_torch_tf32_available() -> bool:  # pragma: no cover
     return not (
         not torch.cuda.is_available()
         or torch.version.cuda is None
-        or torch.cuda.get_device_properties(torch.cuda.current_device()).major
-        < 8
         or int(torch.version.cuda.split(".", maxsplit=1)[0]) < 11
         or version.parse(torch.__version__) < version.parse("1.7")
     )

vis4d/config/common/models/faster_rcnn.py

@@ -38,7 +38,7 @@ def get_default_rpn_box_codec_cfg(
     target_stds: tuple[float, ...] = (1.0, 1.0, 1.0, 1.0),
 ) -> tuple[ConfigDict, ConfigDict]:
     """Get default config for rpn box encoder and decoder."""
-    return tuple(  # type: ignore
+    return tuple(
         class_config(x, target_means=target_means, target_stds=target_stds)
         for x in (DeltaXYWHBBoxEncoder, DeltaXYWHBBoxDecoder)
     )
@@ -49,7 +49,7 @@ def get_default_rcnn_box_codec_cfg(
     target_stds: tuple[float, ...] = (0.1, 0.1, 0.2, 0.2),
 ) -> tuple[ConfigDict, ConfigDict]:
     """Get default config for rcnn box encoder and decoder."""
-    return tuple(  # type: ignore
+    return tuple(
         class_config(x, target_means=target_means, target_stds=target_stds)
         for x in (DeltaXYWHBBoxEncoder, DeltaXYWHBBoxDecoder)
     )

vis4d/data/transforms/pad.py

@@ -58,7 +58,9 @@ def __call__(self, images: list[NDArrayF32]) -> list[NDArrayF32]:
         for i, (image, h, w) in enumerate(zip(images, heights, widths)):
             pad_param = (0, max_hw[1] - w, 0, max_hw[0] - h)
             image_ = torch.from_numpy(image).permute(0, 3, 1, 2)
-            image_ = F.pad(image_, pad_param, self.mode, self.value)
+            image_ = F.pad(  # pylint: disable=not-callable
+                image_, pad_param, self.mode, self.value
+            )
             images[i] = image_.permute(0, 2, 3, 1).numpy()
         return images
 

vis4d/data/transforms/point_sampling.py

@@ -178,7 +178,7 @@ def __call__(
                     selected_idxs = self._idx_sampler([coordinates])[0]
                     if selected_idxs.sum() >= self.min_pts:
                         sampled_idxs.append(selected_idxs)
-        return [np.stack(sampled_idxs)] * len(coordinates_list)
+        return [np.stack(sampled_idxs)] * len(coordinates_list)  # type: ignore
 
 
 @Transform([K.points3d, "transforms.sampling_idxs"], K.points3d)

vis4d/engine/loss_module.py

@@ -3,6 +3,7 @@
 
 from typing import TypedDict, Union
 
+import torch
 from torch import Tensor, nn
 from typing_extensions import NotRequired
 
@@ -72,14 +73,22 @@ class LossModule(nn.Module):
     weighted by the corresponding weight and returned as a dictionary.
     """
 
-    def __init__(self, losses: list[LossDefinition] | LossDefinition) -> None:
+    def __init__(
+        self,
+        losses: list[LossDefinition] | LossDefinition,
+        exclude_attributes: list[str] | None = None,
+    ) -> None:
         """Creates an instance of the class.
 
         Each loss will be called with arguments matching the kwargs of the loss
         function through its connector. By default, the weight is set to 1.0.
 
         Args:
             losses (list[LossDefinition]): List of loss definitions.
+            exclude_attributes (list[str] | None): List of attributes returned
+                by the losses that should be excluded from the total loss
+                computation. Use it to log metrics that should not be
+                optimised. Defaults to None.
 
         Example:
             >>> loss = LossModule(
@@ -127,7 +136,11 @@ def __init__(self, losses: list[LossDefinition] | LossDefinition) -> None:
 
             self.losses.append(loss)
 
-    def forward(self, output: DictData, batch: DictData) -> LossesType:
+        self.exclude_attributes = exclude_attributes
+
+    def forward(
+        self, output: DictData, batch: DictData
+    ) -> tuple[Tensor, dict[str, float]]:
         """Forward of loss module.
 
         This function will call all loss functions and return a dictionary
@@ -142,7 +155,8 @@ def forward(self, output: DictData, batch: DictData) -> LossesType:
             batch (DictData): Batch data.
 
         Returns:
-            LossesType: The loss values.
+            total_loss: The total loss value.
+            metrics: The metrics disctionary.
         """
         loss_dict: LossesType = {}
 
@@ -182,6 +196,26 @@ def forward(self, output: DictData, batch: DictData) -> LossesType:
                 while key in loss_dict:
                     key = "__" + key
 
-                loss_dict[key] = loss_weight * value
+                loss_dict[key] = torch.mul(loss_weight, value)
+
+        # Convert loss_dict to total loss and metrics dictionary
+        metrics: dict[str, float] = {}
+        if isinstance(loss_dict, Tensor):
+            total_loss = loss_dict
+        elif isinstance(loss_dict, dict):
+            keep_loss_dict: LossesType = {}
+            for k, v in loss_dict.items():
+                metrics[k] = v.detach().cpu().item()
+                if (
+                    self.exclude_attributes is None
+                    or k not in self.exclude_attributes
+                ):
+                    keep_loss_dict[k] = v
+            total_loss = sum(keep_loss_dict.values())  # type: ignore
+        else:
+            raise TypeError(
+                "Loss function must return a Tensor or a dict of Tensor"
+            )
+        metrics["loss"] = total_loss.detach().cpu().item()
 
-        return loss_dict
+        return total_loss, metrics

vis4d/engine/trainer.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import torch
-from torch import Tensor, nn
+from torch import nn
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
@@ -236,21 +236,7 @@ def fit(
                 # Forward + backward + optimize
                 output = model(**self.train_data_connector(data))
 
-                losses = loss_module(output, data)
-
-                metrics: dict[str, float] = {}
-                if isinstance(losses, Tensor):
-                    total_loss = losses
-                elif isinstance(losses, dict):
-                    total_loss = sum(losses.values())  # type: ignore
-                    for k, v in losses.items():
-                        metrics[k] = v.detach().cpu().item()
-                else:
-                    raise TypeError(
-                        "Loss function must return a Tensor or a dict of "
-                        + "Tensor"
-                    )
-                metrics["loss"] = total_loss.detach().cpu().item()
+                total_loss, metrics = loss_module(output, data)
 
                 total_loss.backward()
 

vis4d/model/adapter/ema.py

@@ -61,7 +61,7 @@ def update(self, steps: int) -> None:  # pylint: disable=unused-argument
         """Update the internal EMA model."""
         self._update(
             self.model,
-            update_fn=lambda e, m: self.decay * e + (1.0 - self.decay) * m,
+            update_fn=lambda e, m: self.decay * e + (1.0 - self.decay) * m,  # type: ignore # pylint: disable=line-too-long
         )
 
     def set(self, model: nn.Module) -> None:
@@ -113,5 +113,5 @@ def update(self, steps: int) -> None:
         )
         self._update(
             self.model,
-            update_fn=lambda e, m: decay * e + (1.0 - decay) * m,
+            update_fn=lambda e, m: decay * e + (1.0 - decay) * m,  # type: ignore # pylint: disable=line-too-long
         )

vis4d/model/detect/mask_rcnn.py

@@ -180,7 +180,7 @@ def forward_train(
         assert outputs.sampled_proposals is not None
         assert outputs.sampled_targets is not None
         pos_proposals = apply_mask(
-            [label == 1 for label in outputs.sampled_targets.labels],
+            [torch.eq(label, 1) for label in outputs.sampled_targets.labels],
             outputs.sampled_proposals.boxes,
         )[0]
         mask_outs = self.mask_head(features, pos_proposals)