first commit

dnnultis/log/__init__.py

@@ -0,0 +1,2 @@
+from .logger import *
+from .wandb import *

dnnultis/log/logger.py

@@ -0,0 +1,86 @@
+import functools
+import logging
+import os, sys
+from termcolor import colored
+
+
+class _ColorfulFormatter(logging.Formatter):
+    def __init__(self, *args, **kwargs):
+        self._root_name = kwargs.pop("root_name") + "."
+        self._abbrev_name = kwargs.pop("abbrev_name", "")
+        if len(self._abbrev_name):
+            self._abbrev_name = self._abbrev_name + "."
+        super(_ColorfulFormatter, self).__init__(*args, **kwargs)
+
+    def formatMessage(self, record):
+        record.name = record.name.replace(self._root_name, self._abbrev_name)
+        log = super(_ColorfulFormatter, self).formatMessage(record)
+        if record.levelno == logging.WARNING:
+            prefix = colored("WARNING", "red", attrs=["blink"])
+        elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:
+            prefix = colored("ERROR", "red", attrs=["blink", "underline"])
+        else:
+            return log
+        return prefix + " " + log
+
+
+# so that calling setup_logging.root multiple times won't add many handlers
+@functools.lru_cache()
+def setup_logging(output=None,
+                      distributed_rank=0,
+                      default_level=logging.INFO,
+                      *,
+                      color=True,
+                      name=__name__):
+    """
+    Initialize the detectron2 logging.root and set its verbosity level to "INFO".
+    Args:
+        output (str): a file name or a directory to save log. If None, will not save log file.
+            If ends with ".txt" or ".log", assumed to be a file name.
+            Otherwise, logs will be saved to `output/log.txt`.
+        name (str): the root module name of this logging.root
+    Returns:
+        logging.logging.root: a logging.root
+    """
+    logging.root.setLevel(default_level)
+    logging.root.propagate = False
+
+    plain_formatter = logging.Formatter(
+        "[%(asctime)s] %(name)s %(levelname)s: %(message)s",
+        datefmt="%y/%m/%d %H:%M:%S")
+    # stdout logging: master only
+    if distributed_rank == 0:
+        ch = logging.StreamHandler(stream=sys.stdout)
+        ch.setLevel(default_level)
+        if color:
+            formatter = _ColorfulFormatter(
+                colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s",
+                datefmt="%y/%m/%d %H:%M:%S",
+                root_name=name,
+            )
+        else:
+            formatter = plain_formatter
+        ch.setFormatter(formatter)
+        logging.root.addHandler(ch)
+
+    # file logging: all workers
+    if output is not None:
+        if output.endswith(".txt") or output.endswith(".log"):
+            filename = output
+        else:
+            filename = os.path.join(output, "log.txt")
+        if distributed_rank > 0:
+            filename = filename + f".rank{distributed_rank}"
+        os.makedirs(os.path.dirname(filename), exist_ok=True)
+
+        fh = logging.StreamHandler(_cached_log_stream(filename))
+        fh.setLevel(default_level)
+        fh.setFormatter(plain_formatter)
+        logging.root.addHandler(fh)
+
+
+# cache the opened file object, so that different calls to `setup_logging.root`
+# with the same file name can safely write to the same file.
+@functools.lru_cache(maxsize=None)
+def _cached_log_stream(filename):
+    return open(filename, "a")

dnnultis/log/wandb.py

@@ -0,0 +1,84 @@
+import shutil
+import os
+import subprocess
+import wandb
+
+
+class WandbUrls:
+    def __init__(self, url):
+
+        hash = url.split("/")[-2]
+        project = url.split("/")[-3]
+        entity = url.split("/")[-4]
+
+        self.weight_url = url
+        self.log_url = "https://app.wandb.ai/{}/{}/runs/{}/logs".format(entity, project, hash)
+        self.chart_url = "https://app.wandb.ai/{}/{}/runs/{}".format(entity, project, hash)
+        self.overview_url = "https://app.wandb.ai/{}/{}/runs/{}/overview".format(entity, project, hash)
+        self.config_url = "https://app.wandb.ai/{}/{}/runs/{}/files/hydra-config.yaml".format(
+            entity, project, hash
+        )
+        self.overrides_url = "https://app.wandb.ai/{}/{}/runs/{}/files/overrides.yaml".format(entity, project, hash)
+
+    def __repr__(self):
+        msg = "=================================================== WANDB URLS ===================================================================\n"
+        for k, v in self.__dict__.items():
+            msg += "{}: {}\n".format(k.upper(), v)
+        msg += "=================================================================================================================================\n"
+        return msg
+
+
+class Wandb:
+    IS_ACTIVE = False
+
+    @staticmethod
+    def set_urls_to_model(model, url):
+        wandb_urls = WandbUrls(url)
+        model.wandb = wandb_urls
+
+    @staticmethod
+    def _set_to_wandb_args(wandb_args, cfg, name):
+        var = getattr(cfg.wandb, name, None)
+        if var:
+            wandb_args[name] = var
+
+    @staticmethod
+    def launch(cfg, launch: bool):
+        if launch:
+
+            Wandb.IS_ACTIVE = True
+
+            wandb_args = {}
+            wandb_args["resume"] = "allow"
+            Wandb._set_to_wandb_args(wandb_args, cfg, "tags")
+            Wandb._set_to_wandb_args(wandb_args, cfg, "project")
+            Wandb._set_to_wandb_args(wandb_args, cfg, "name")
+            Wandb._set_to_wandb_args(wandb_args, cfg, "entity")
+            Wandb._set_to_wandb_args(wandb_args, cfg, "notes")
+            Wandb._set_to_wandb_args(wandb_args, cfg, "config")
+            Wandb._set_to_wandb_args(wandb_args, cfg, "id")
+
+            try:
+                commit_sha = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("ascii").strip()
+                gitdiff = subprocess.check_output(["git", "diff", "--", "':!notebooks'"]).decode()
+            except:
+                commit_sha = "n/a"
+                gitdiff = ""
+
+            config = wandb_args.get("config", {})
+            wandb_args["config"] = {
+                **config,
+                "run_path": os.getcwd(),
+                "commit": commit_sha,
+                "gitdiff": gitdiff
+            }
+            wandb.init(**wandb_args, sync_tensorboard=True)
+            wandb.save(os.path.join(os.getcwd(), cfg.cfg_path))
+
+    @staticmethod
+    def add_file(file_path: str):
+        if not Wandb.IS_ACTIVE:
+            raise RuntimeError("wandb is inactive, please launch first.")
+
+        filename = os.path.basename(file_path)
+        shutil.copyfile(file_path, os.path.join(wandb.run.dir, filename))