diff --git a/micromind/core.py b/micromind/core.py index 047acbe..4d149d0 100644 --- a/micromind/core.py +++ b/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,7 +499,8 @@ 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)) @@ -485,7 +508,11 @@ def validate(self) -> Dict: 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() @@ -493,10 +520,7 @@ def validate(self) -> Dict: 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 " diff --git a/micromind/utils/yolo_helpers.py b/micromind/utils/yolo_helpers.py index b58dfc3..2258aee 100644 --- a/micromind/utils/yolo_helpers.py +++ b/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, diff --git a/recipes/objection_detection/extra_requirements.txt b/recipes/objection_detection/extra_requirements.txt new file mode 100644 index 0000000..8a371b3 --- /dev/null +++ b/recipes/objection_detection/extra_requirements.txt @@ -0,0 +1,2 @@ +opencv-python +ultralytics==8.0.215 diff --git a/recipes/objection_detection/train.py b/recipes/objection_detection/train.py index be1dc22..33f7523 100644 --- a/recipes/objection_detection/train.py +++ b/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,