| """This module implements the writer class for logging to tensorboard or wandb.""" |
|
|
| import logging |
| import os |
| from typing import Any, Dict, Optional |
|
|
| from omegaconf import DictConfig |
| from torch import nn |
| from torch.utils.tensorboard import SummaryWriter as TFSummaryWriter |
|
|
| from siclib import __module_name__ |
|
|
| logger = logging.getLogger(__name__) |
|
|
| try: |
| import wandb |
| except ImportError: |
| logger.debug("Could not import wandb.") |
| wandb = None |
|
|
| |
|
|
|
|
| def dot_conf(conf: DictConfig) -> Dict[str, Any]: |
| """Recursively convert a DictConfig to a flat dict with keys joined by dots.""" |
| d = {} |
| for k, v in conf.items(): |
| if isinstance(v, DictConfig): |
| d |= {f"{k}.{k2}": v2 for k2, v2 in dot_conf(v).items()} |
| else: |
| d[k] = v |
| return d |
|
|
|
|
| class SummaryWriter: |
| """Writer class for logging to tensorboard or wandb.""" |
|
|
| def __init__(self, conf: DictConfig, args: DictConfig, log_dir: str): |
| """Initialize the writer.""" |
| self.conf = conf |
|
|
| if not conf.train.writer: |
| self.use_wandb = False |
| self.use_tensorboard = False |
| return |
|
|
| self.use_wandb = "wandb" in conf.train.writer |
| self.use_tensorboard = "tensorboard" in conf.train.writer |
|
|
| if self.use_wandb and not wandb: |
| raise ImportError("wandb not installed.") |
|
|
| if self.use_tensorboard: |
| self.writer = TFSummaryWriter(log_dir=log_dir) |
|
|
| if self.use_wandb: |
| os.environ["WANDB__SERVICE_WAIT"] = "300" |
| wandb.init(project=__module_name__, name=args.experiment, config=dot_conf(conf)) |
|
|
| if conf.train.writer and not self.use_wandb and not self.use_tensorboard: |
| raise NotImplementedError(f"Writer {conf.train.writer} not implemented") |
|
|
| def add_scalar(self, tag: str, value: float, step: Optional[int] = None): |
| """Log a scalar value to tensorboard or wandb.""" |
| if self.use_wandb: |
| step = 1 if step == 0 else step |
| wandb.log({tag: value}, step=step) |
|
|
| if self.use_tensorboard: |
| self.writer.add_scalar(tag, value, step) |
|
|
| def add_figure(self, tag: str, figure, step: Optional[int] = None): |
| """Log a figure to tensorboard or wandb.""" |
| if self.use_wandb: |
| step = 1 if step == 0 else step |
| wandb.log({tag: figure}, step=step) |
| if self.use_tensorboard: |
| self.writer.add_figure(tag, figure, step) |
|
|
| def add_histogram(self, tag: str, values, step: Optional[int] = None): |
| """Log a histogram to tensorboard or wandb.""" |
| if self.use_tensorboard: |
| self.writer.add_histogram(tag, values, step) |
|
|
| def add_text(self, tag: str, text: str, step: Optional[int] = None): |
| """Log text to tensorboard or wandb.""" |
| if self.use_tensorboard: |
| self.writer.add_text(tag, text, step) |
|
|
| def add_pr_curve(self, tag: str, values, step: Optional[int] = None): |
| """Log a precision-recall curve to tensorboard or wandb.""" |
| if self.use_wandb: |
| step = 1 if step == 0 else step |
| |
| |
| wandb.log({tag: wandb.plots.precision_recall(values)}, step=step) |
|
|
| if self.use_tensorboard: |
| self.writer.add_pr_curve(tag, values, step) |
|
|
| def watch(self, model: nn.Module, log_freq: int = 1000): |
| """Watch a model for gradient updates.""" |
| if self.use_wandb: |
| wandb.watch( |
| model, |
| log="gradients", |
| log_freq=log_freq, |
| ) |
|
|
| def close(self): |
| """Close the writer.""" |
| if self.use_wandb: |
| wandb.finish() |
|
|
| if self.use_tensorboard: |
| self.writer.close() |
|
|
