Fix and improvements to the Metric class

micromind/core.py

@@ -5,7 +5,7 @@
 Authors:
     - Francesco Paissan, 2023
 """
-from typing import Dict, Union, Tuple, Callable, List
+from typing import Dict, Union, Tuple, Callable, List, Optional
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from argparse import Namespace
@@ -82,15 +82,23 @@ class Metric:
         0.5
     """
 
-    def __init__(self, name: str, fn: Callable, reduction="mean"):
+    def __init__(
+        self,
+        name: str,
+        fn: Callable,
+        reduction: Optional[str] = "mean",
+        eval_only: Optional[bool] = False,
+        eval_period: Optional[int] = 1,
+    ):
         self.name = name
         self.fn = fn
         self.reduction = reduction
+        self.eval_only = eval_only
+        self.eval_period = eval_period
+
         self.history = {s: [] for s in [Stage.train, Stage.val, Stage.test]}
 
     def __call__(self, pred, batch, stage, device="cpu"):
-        # if pred.device != device:
-        #     pred = pred.to(device)
         dat = self.fn(pred, batch)
         if dat.ndim == 0:
             dat = dat.unsqueeze(0)
@@ -341,6 +349,9 @@ def on_train_end(self):
             logger.info(f"Removed temporary folder {self.experiment_folder}.")
             shutil.rmtree(self.experiment_folder)
 
+    def eval(self):
+        self.modules.eval()
+
     def train(
         self,
         epochs: int = 1,
@@ -385,6 +396,7 @@ def train(
             )
         with self.accelerator.autocast():
             for e in range(self.start_epoch, epochs):
+                self.current_epoch = e
                 pbar = tqdm(
                     self.datasets["train"],
                     unit="batches",
@@ -409,11 +421,17 @@ def train(
                     self.opt.step()
 
                     for m in self.metrics:
-                        m(model_out, batch, Stage.train, self.device)
+                        if (
+                            self.current_epoch + 1
+                        ) % m.eval_period == 0 and not m.eval_only:
+                            m(model_out, batch, Stage.train, self.device)
 
-                    running_train = {
-                        "train_" + m.name: m.reduce(Stage.train) for m in self.metrics
-                    }
+                    running_train = {}
+                    for m in self.metrics:
+                        if (
+                            self.current_epoch + 1
+                        ) % m.eval_period == 0 and not m.eval_only:
+                            running_train["train_" + m.name] = m.reduce(Stage.train)
 
                     running_train.update({"train_loss": loss_epoch / (idx + 1)})
 
@@ -424,9 +442,13 @@ def train(
 
                 pbar.close()
 
-                train_metrics = {
-                    "train_" + m.name: m.reduce(Stage.train, True) for m in self.metrics
-                }
+                train_metrics = {}
+                for m in self.metrics:
+                    if (
+                        self.current_epoch + 1
+                    ) % m.eval_period == 0 and not m.eval_only:
+                        train_metrics["train_" + m.name] = m.reduce(Stage.train, True)
+
                 train_metrics.update({"train_loss": loss_epoch / (idx + 1)})
 
                 if "val" in datasets:
@@ -477,26 +499,28 @@ def validate(self) -> Dict:
                 model_out = self(batch)
                 loss = self.compute_loss(model_out, batch)
                 for m in self.metrics:
-                    m(model_out, batch, Stage.val, self.device)
+                    if (self.current_epoch + 1) % m.eval_period == 0:
+                        m(model_out, batch, Stage.val, self.device)
 
                 loss_epoch += loss.item()
                 pbar.set_postfix(loss=loss_epoch / (idx + 1))
 
                 if self.debug and idx > 10:
                     break
 
-        val_metrics = {"val_" + m.name: m.reduce(Stage.val, True) for m in self.metrics}
+        val_metrics = {}
+        for m in self.metrics:
+            if (self.current_epoch + 1) % m.eval_period == 0:
+                val_metrics = {"val_" + m.name: m.reduce(Stage.val, True)}
+
         val_metrics.update({"val_loss": loss_epoch / (idx + 1)})
 
         pbar.close()
 
         return val_metrics
 
     @torch.no_grad()
-    def test(
-        self,
-        datasets: Dict = {},
-        metrics: List[Metric] = []) -> None:
+    def test(self, datasets: Dict = {}, metrics: List[Metric] = []) -> None:
         """Runs the test steps."""
         assert "test" in datasets, "Test dataloader was not specified."
         self.modules.eval()
@@ -525,9 +549,7 @@ def test(
 
         pbar.close()
 
-        test_metrics = {
-            "test_" + m.name: m.reduce(Stage.test, True) for m in metrics
-        }
+        test_metrics = {"test_" + m.name: m.reduce(Stage.test, True) for m in metrics}
         test_metrics.update({"test_loss": loss_epoch / (idx + 1)})
         s_out = (
             "Testing "

micromind/utils/yolo_helpers.py

@@ -57,7 +57,7 @@ def load_config(file_path):
             "path": path,
             "train": train.as_posix(),
             "val": val.as_posix(),
-            "test": test.as_posix(),
+            "test": test,
             "names": config["names"],
             "download": config.get("download"),
             "yaml_file": file_path,

recipes/objection_detection/extra_requirements.txt

@@ -0,0 +1,2 @@
+opencv-python
+ultralytics==8.0.215

recipes/objection_detection/train.py

@@ -255,12 +255,12 @@ def mAP(self, pred, batch):
 if __name__ == "__main__":
     batch_size = 8
 
-    m_cfg, data_cfg = load_config("cfg/coco.yaml")
+    m_cfg, data_cfg = load_config("cfg/coco8.yaml")
 
     mode = "train"
     coco8_dataset = build_yolo_dataset(
         m_cfg,
-        "datasets/coco/images/train2017",
+        "datasets/coco8/images/train",
         batch_size,
         data_cfg,
         mode=mode,
@@ -278,7 +278,7 @@ def mAP(self, pred, batch):
     mode = "val"
     coco8_dataset = build_yolo_dataset(
         m_cfg,
-        "datasets/coco/images/val2017",
+        "datasets/coco8/images/val",
         batch_size,
         data_cfg,
         mode=mode,
@@ -295,10 +295,10 @@ def mAP(self, pred, batch):
 
     hparams = parse_arguments()
     m = YOLO(m_cfg, hparams=hparams)
-    mAP = Metric("mAP", m.mAP)
+    mAP = Metric("mAP", m.mAP, eval_only=False, eval_period=2)
 
     m.train(
-        epochs=50,
+        epochs=2,
         datasets={"train": train_loader, "val": val_loader},
         metrics=[mAP],
         debug=hparams.debug,