From 4383f1ead609b4d573ba37f2d7391249af708c4c Mon Sep 17 00:00:00 2001 From: sfmig <33267254+sfmig@users.noreply.github.com> Date: Tue, 5 Nov 2024 20:53:50 +0000 Subject: [PATCH 1/5] Add mlflow skeleton --- crabs/tracker/track_video.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py index 7971ef3e..fa66e8ee 100644 --- a/crabs/tracker/track_video.py +++ b/crabs/tracker/track_video.py @@ -13,6 +13,10 @@ import yaml # type: ignore from crabs.detector.models import FasterRCNN +from crabs.detector.utils.detection import ( + log_mlflow_metadata_as_info, + set_mlflow_run_name, +) from crabs.tracker.evaluate_tracker import TrackerEvaluate from crabs.tracker.sort import Sort from crabs.tracker.utils.io import ( @@ -61,6 +65,14 @@ def __init__(self, args: argparse.Namespace) -> None: iou_threshold=self.config["iou_threshold"], ) + # MLflow experiment name and run name + self.experiment_name = args.experiment_name + self.run_name = set_mlflow_run_name() + self.mlflow_folder = args.mlflow_folder + + # Log MLflow information to screen + log_mlflow_metadata_as_info(self) + def setup(self): """Load tracking config, trained model and input video path.""" with open(self.config_file) as f: @@ -319,6 +331,24 @@ def tracking_parse_args(args): "Valid inputs are: cpu or gpu. Default: gpu." ), ) + parser.add_argument( + "--experiment_name", + type=str, + help=( + "Name of the experiment in MLflow, under which the current run " + "will be logged. " + "By default: _evaluation." + ), + ) + parser.add_argument( + "--mlflow_folder", + type=str, + default="./ml-runs", + help=( + "Path to MLflow directory where to log the evaluation data. " + "Default: 'ml-runs' directory under the current working directory." + ), + ) parser.add_argument( "--max_frames_to_read", type=int, From af381bbbf5c9f882c3eba4a246857e4adf87e39c Mon Sep 17 00:00:00 2001 From: sfmig <33267254+sfmig@users.noreply.github.com> Date: Tue, 5 Nov 2024 21:09:19 +0000 Subject: [PATCH 2/5] Add skeleton for predict_step --- crabs/detector/models.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/crabs/detector/models.py b/crabs/detector/models.py index a6bc8a4b..4b1f93c0 100644 --- a/crabs/detector/models.py +++ b/crabs/detector/models.py @@ -243,3 +243,37 @@ def configure_optimizers(self) -> dict[str, torch.optim.Optimizer]: weight_decay=self.config["wdecay"], ) return {"optimizer": optimizer} + + # def predict_step( + # self, images: torch.Tensor, dataloader_idx=0 + # ) -> torch.Tensor: + # """Perform inference on a batch of images.""" + + # # ensure batch is on device + # # batch = batch.to(self.device) + # # torch.Size([2160, 4096, 3]) + + # # apply inference transform + # # test_val_transforms = CrabsDataModule._get_test_val_transform(None) + # # inference_transforms = transforms.Compose( + # # [ + # # transforms.ToImage(), + # # transforms.ToDtype(torch.float32, scale=True), + # # ] + # # ) + # # batch_transformed = inference_transforms(batch) # B, C, H, W + # # breakpoint() + + # # # run prediction + # # predictions = self.model(images) + # transform = transforms.Compose( + # [ + # transforms.ToImage(), + # transforms.ToDtype(torch.float32, scale=True), + # ] + # ) + # image_tensors = transform(images).to(self.device) + # # image_tensors = image_tensors.unsqueeze(0) + # with torch.no_grad(): + # predictions = self.model(image_tensors) + # return predictions From abaeb4652c08cb62e121daa5a6ce55a7bed53140 Mon Sep 17 00:00:00 2001 From: sfmig <33267254+sfmig@users.noreply.github.com> Date: Tue, 5 Nov 2024 21:17:47 +0000 Subject: [PATCH 3/5] Fix test --- tests/test_unit/test_track_video.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_unit/test_track_video.py b/tests/test_unit/test_track_video.py index 7d7ffa89..31f5d92c 100644 --- a/tests/test_unit/test_track_video.py +++ b/tests/test_unit/test_track_video.py @@ -20,6 +20,8 @@ def mock_args(): accelerator="gpu", annotations_file=None, save_video=None, + experiment_name=None, + mlflow_folder=Path(temp_dir) / "ml-runs", ) From 18ad17fb50d5d50b1db4da25daefc140c1b2274d Mon Sep 17 00:00:00 2001 From: sfmig <33267254+sfmig@users.noreply.github.com> Date: Tue, 5 Nov 2024 21:25:17 +0000 Subject: [PATCH 4/5] Add setup_trainer --- crabs/tracker/track_video.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py index fa66e8ee..41497041 100644 --- a/crabs/tracker/track_video.py +++ b/crabs/tracker/track_video.py @@ -7,6 +7,7 @@ from pathlib import Path import cv2 +import lightning import numpy as np import torch import torchvision.transforms.v2 as transforms @@ -16,6 +17,7 @@ from crabs.detector.utils.detection import ( log_mlflow_metadata_as_info, set_mlflow_run_name, + setup_mlflow_logger, ) from crabs.tracker.evaluate_tracker import TrackerEvaluate from crabs.tracker.sort import Sort @@ -73,6 +75,33 @@ def __init__(self, args: argparse.Namespace) -> None: # Log MLflow information to screen log_mlflow_metadata_as_info(self) + def setup_trainer(self): + """Set up trainer object with logging for testing.""" + # Setup logger + mlf_logger = setup_mlflow_logger( + experiment_name=self.experiment_name, + run_name=self.run_name, + mlflow_folder=self.mlflow_folder, + cli_args=self.args, + ) + + # Add trained model section to MLflow hyperparameters + mlf_logger.log_hyperparams( + { + "trained_model/experiment_name": self.trained_model_expt_name, + "trained_model/run_name": self.trained_model_run_name, + "trained_model/ckpt_file": Path(self.trained_model_path).name, + } + ) + + # Add other unlogged information from init? + + # Return trainer linked to logger + return lightning.Trainer( + accelerator=self.accelerator, # lightning accelerators + logger=mlf_logger, + ) + def setup(self): """Load tracking config, trained model and input video path.""" with open(self.config_file) as f: @@ -334,10 +363,11 @@ def tracking_parse_args(args): parser.add_argument( "--experiment_name", type=str, + default="Inference", help=( "Name of the experiment in MLflow, under which the current run " "will be logged. " - "By default: _evaluation." + "By default: Inference." ), ) parser.add_argument( From c3758342a306d5780d33a651358b085fa26a09e7 Mon Sep 17 00:00:00 2001 From: sfmig <33267254+sfmig@users.noreply.github.com> Date: Tue, 5 Nov 2024 21:32:51 +0000 Subject: [PATCH 5/5] Add slurm logs as artifacts --- crabs/tracker/track_video.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py index 41497041..b3d4efcf 100644 --- a/crabs/tracker/track_video.py +++ b/crabs/tracker/track_video.py @@ -270,6 +270,12 @@ def run_tracking(self): if self.args.save_frames: close_csv_file(self.csv_file) + # if this is a slurm job: add slurm logs as artifacts + # slurm_job_id = os.environ.get("SLURM_JOB_ID") + # slurm_job_name = os.environ.get("SLURM_JOB_NAME") + # if slurm_job_id and (slurm_job_name != "bash"): + # slurm_logs_as_artifacts(trainer.logger, slurm_job_id) + def main(args) -> None: """Run detection+tracking inference on video.