Comet Logging and Visualization Integration (#9232)

* add comet to logger interface * add comet logger * add support for updated parameters * clean up offline logger creation * update callback args for comet logger * add comet optimizer * add optimizer config * add comet README * update tutorial notebook with Comet section * add option to log class level metrics * add support for class level metrics and confusion matrix * handle errors when adding files to artifacts * fix typo * clean resume workflow * updates for HPO * update comet README * fix typo in comet README * update code snippets in comet README * update comet links in tutorial * updated links * change optimizer batch size param and update comet README image * update comet section in tutorial * use prexisting cmd line flags to configure logger * update artifact upload/download flow * remove come remove comet logger specific cmd line args * move downloading weights into comet logger code * remove extra argparse * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * change checkpoint logging flow to follow offline logger * update resume flow * add comet logger to remote dataset property * update cmd line args in hpo * set types for integer/float env variables * update README * fix typo in README * default to always logging model predictions * Update tutorial.ipynb * Update train.py * Add Comet to Integrations table * Update README.md * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2022-09-07 18:28:46 +03:00
1 parent 5a134e0653
commit 903b239f13
@@ -0,0 +1,150 @@
+import logging
+import os
+from urllib.parse import urlparse
+
+try:
+    import comet_ml
+except (ModuleNotFoundError, ImportError):
+    comet_ml = None
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+COMET_PREFIX = "comet://"
+COMET_MODEL_NAME = os.getenv("COMET_MODEL_NAME", "yolov5")
+COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv("COMET_DEFAULT_CHECKPOINT_FILENAME", "last.pt")
+
+
+def download_model_checkpoint(opt, experiment):
+    model_dir = f"{opt.project}/{experiment.name}"
+    os.makedirs(model_dir, exist_ok=True)
+
+    model_name = COMET_MODEL_NAME
+    model_asset_list = experiment.get_model_asset_list(model_name)
+
+    if len(model_asset_list) == 0:
+        logger.error(f"COMET ERROR: No checkpoints found for model name : {model_name}")
+        return
+
+    model_asset_list = sorted(
+        model_asset_list,
+        key=lambda x: x["step"],
+        reverse=True,
+    )
+    logged_checkpoint_map = {asset["fileName"]: asset["assetId"] for asset in model_asset_list}
+
+    resource_url = urlparse(opt.weights)
+    checkpoint_filename = resource_url.query
+
+    if checkpoint_filename:
+        asset_id = logged_checkpoint_map.get(checkpoint_filename)
+    else:
+        asset_id = logged_checkpoint_map.get(COMET_DEFAULT_CHECKPOINT_FILENAME)
+        checkpoint_filename = COMET_DEFAULT_CHECKPOINT_FILENAME
+
+    if asset_id is None:
+        logger.error(f"COMET ERROR: Checkpoint {checkpoint_filename} not found in the given Experiment")
+        return
+
+    try:
+        logger.info(f"COMET INFO: Downloading checkpoint {checkpoint_filename}")
+        asset_filename = checkpoint_filename
+
+        model_binary = experiment.get_asset(asset_id, return_type="binary", stream=False)
+        model_download_path = f"{model_dir}/{asset_filename}"
+        with open(model_download_path, "wb") as f:
+            f.write(model_binary)
+
+        opt.weights = model_download_path
+
+    except Exception as e:
+        logger.warning("COMET WARNING: Unable to download checkpoint from Comet")
+        logger.exception(e)
+
+
+def set_opt_parameters(opt, experiment):
+    """Update the opts Namespace with parameters
+    from Comet's ExistingExperiment when resuming a run
+
+    Args:
+        opt (argparse.Namespace): Namespace of command line options
+        experiment (comet_ml.APIExperiment): Comet API Experiment object
+    """
+    asset_list = experiment.get_asset_list()
+    resume_string = opt.resume
+
+    for asset in asset_list:
+        if asset["fileName"] == "opt.yaml":
+            asset_id = asset["assetId"]
+            asset_binary = experiment.get_asset(asset_id, return_type="binary", stream=False)
+            opt_dict = yaml.safe_load(asset_binary)
+            for key, value in opt_dict.items():
+                setattr(opt, key, value)
+            opt.resume = resume_string
+
+    # Save hyperparameters to YAML file
+    # Necessary to pass checks in training script
+    save_dir = f"{opt.project}/{experiment.name}"
+    os.makedirs(save_dir, exist_ok=True)
+
+    hyp_yaml_path = f"{save_dir}/hyp.yaml"
+    with open(hyp_yaml_path, "w") as f:
+        yaml.dump(opt.hyp, f)
+    opt.hyp = hyp_yaml_path
+
+
+def check_comet_weights(opt):
+    """Downloads model weights from Comet and updates the
+    weights path to point to saved weights location
+
+    Args:
+        opt (argparse.Namespace): Command Line arguments passed
+            to YOLOv5 training script
+
+    Returns:
+        None/bool: Return True if weights are successfully downloaded
+            else return None
+    """
+    if comet_ml is None:
+        return
+
+    if isinstance(opt.weights, str):
+        if opt.weights.startswith(COMET_PREFIX):
+            api = comet_ml.API()
+            resource = urlparse(opt.weights)
+            experiment_path = f"{resource.netloc}{resource.path}"
+            experiment = api.get(experiment_path)
+            download_model_checkpoint(opt, experiment)
+            return True
+
+    return None
+
+
+def check_comet_resume(opt):
+    """Restores run parameters to its original state based on the model checkpoint
+    and logged Experiment parameters.
+
+    Args:
+        opt (argparse.Namespace): Command Line arguments passed
+            to YOLOv5 training script
+
+    Returns:
+        None/bool: Return True if the run is restored successfully
+            else return None
+    """
+    if comet_ml is None:
+        return
+
+    if isinstance(opt.resume, str):
+        if opt.resume.startswith(COMET_PREFIX):
+            api = comet_ml.API()
+            resource = urlparse(opt.resume)
+            experiment_path = f"{resource.netloc}{resource.path}"
+            experiment = api.get(experiment_path)
+            set_opt_parameters(opt, experiment)
+            download_model_checkpoint(opt, experiment)
+
+            return True
+
+    return None