WandBLogger

minnt.loggers.WandBLogger

Bases: BaseLogger

A W&B logger interface.

The text values are by default also logged as HTML for better visualization.

Source code in minnt/loggers/wandb_logger.py

class WandBLogger(BaseLogger):
    """A W&B logger interface.

    The text values are by default also logged as HTML for better visualization.
    """
    def __init__(self, project: str, *, text_also_as_html: bool = True, **kwargs: Any) -> None:
        """Create the WandBLogger with the given project name.

        Additional keyword arguments are passed to `wandb.init()`.

        Parameters:
          project: The name of the W&B project.
          text_also_as_html: Whether to log text messages also as HTML.
            That has the advantage of interactive visualization of the value
            at different epochs and preserving whitespace formatting.
          kwargs: Additional keyword arguments passed to `wandb.init()`, for example:

            - `dir` ([`str`][str]) – The directory where the W&B files will be stored in a `wandb` subdirectory.
            - `id` ([`str`][str]) – A unique identifier for this run.
            - `name` ([`str`][str]) – A short display name for this run to use in the UI.
            - `notes` ([`str`][str]) – A detailed description of the run.
            - `tags` ([`Sequence`][collections.abc.Sequence][[`str`][str]]) – A list of tags
              to label this run in the UI.
        """
        import wandb
        self.wandb = wandb
        self.run = self.wandb.init(project=project, **kwargs)
        self._text_also_as_html = text_also_as_html

    def close(self) -> None:
        if self.run:
            self.run.finish()
            self.run = None

    def _maybe_as_html(self, label: str, text: str) -> dict[str, Any]:
        """Return a dict with the HTML version of the text if enabled.

        The text is converted to HTML-safe format and returned as a wandb.Html object.
        """
        if not self._text_also_as_html:
            return {}
        return {f"{label}_html": self.wandb.Html("<pre>" + html.escape(text) + "</pre>")}

    def log_audio(self, label: str, audio: AnyArray, sample_rate: int, epoch: int) -> Self:
        audio = self.preprocess_audio(audio).numpy()
        self.run.log({label: self.wandb.Audio(audio, sample_rate=sample_rate)}, step=epoch)
        return self

    def log_config(self, config: dict[str, Any], epoch: int) -> Self:
        self.run.config.update(dict(sorted(config.items())))
        config_json = self.format_config_as_json(config)
        self.run.log({"config": config_json} | self._maybe_as_html("config", config_json), step=epoch)
        return self

    def log_figure(self, label: str, figure: Any, epoch: int, tight_layout: bool = True, close: bool = True) -> Self:
        return super().log_figure(label, figure, epoch, tight_layout, close)

    def log_graph(self, graph: torch.nn.Module, data: TensorOrTensors, epoch: int) -> Self:
        # The logging in WandB has a lot of limitations. One is that all
        # children of the graph must be executed before the graph is logged;
        # however, metrics are usually not executed during the forward pass.
        # We therefore try to exclude metrics from the logged graph.
        watched_children = torch.nn.Module()
        for name, child in graph.named_children():
            if not any(isinstance(module, Metric) for module in child.modules()):
                watched_children.add_module(name, child)
        self.run.watch(watched_children, log=None, log_graph=True)
        with self.graph_in_eval_mode(graph), torch.no_grad():
            graph(*data) if isinstance(data, (list, tuple)) else graph(data)  # Run the graph to log it.
        self.run.unwatch(watched_children)
        return self

    def log_image(self, label: str, image: AnyArray, epoch: int, data_format: DataFormat = "HWC") -> Self:
        image = self.preprocess_image(image, data_format).numpy()
        self.run.log({label: self.wandb.Image(image)}, step=epoch)
        return self

    def log_metrics(self, logs: dict[str, float], epoch: int, description: str | None = None) -> Self:
        self.run.log(logs, step=epoch)
        return self

    def log_text(self, label: str, text: str, epoch: int) -> Self:
        self.run.log({label: text} | self._maybe_as_html(label, text), step=epoch)
        return self

__init__

__init__(
    project: str, *, text_also_as_html: bool = True, **kwargs: Any
) -> None

Create the WandBLogger with the given project name.

Additional keyword arguments are passed to wandb.init().

Parameters:

• project (str) –

  The name of the W&B project.

• text_also_as_html (bool, default: True ) –

  Whether to log text messages also as HTML. That has the advantage of interactive visualization of the value at different epochs and preserving whitespace formatting.

• kwargs (Any, default: {} ) –

  Additional keyword arguments passed to wandb.init(), for example:

  • dir (str) – The directory where the W&B files will be stored in a wandb subdirectory.
  • id (str) – A unique identifier for this run.
  • name (str) – A short display name for this run to use in the UI.
  • notes (str) – A detailed description of the run.
  • tags (Sequence[str]) – A list of tags to label this run in the UI.

Source code in minnt/loggers/wandb_logger.py

def __init__(self, project: str, *, text_also_as_html: bool = True, **kwargs: Any) -> None:
    """Create the WandBLogger with the given project name.

    Additional keyword arguments are passed to `wandb.init()`.

    Parameters:
      project: The name of the W&B project.
      text_also_as_html: Whether to log text messages also as HTML.
        That has the advantage of interactive visualization of the value
        at different epochs and preserving whitespace formatting.
      kwargs: Additional keyword arguments passed to `wandb.init()`, for example:

        - `dir` ([`str`][str]) – The directory where the W&B files will be stored in a `wandb` subdirectory.
        - `id` ([`str`][str]) – A unique identifier for this run.
        - `name` ([`str`][str]) – A short display name for this run to use in the UI.
        - `notes` ([`str`][str]) – A detailed description of the run.
        - `tags` ([`Sequence`][collections.abc.Sequence][[`str`][str]]) – A list of tags
          to label this run in the UI.
    """
    import wandb
    self.wandb = wandb
    self.run = self.wandb.init(project=project, **kwargs)
    self._text_also_as_html = text_also_as_html

close

close() -> None

Close the logger and release its resources.

Source code in minnt/loggers/wandb_logger.py

def close(self) -> None:
    if self.run:
        self.run.finish()
        self.run = None

log_audio

log_audio(label: str, audio: AnyArray, sample_rate: int, epoch: int) -> Self

Log the given audio with the given label at the given epoch.

Parameters:

• label (str) –

  The label of the logged audio.

• audio (AnyArray) –

  The audio to log, represented as an array with any of the following shapes:

  • (L,) of (L, 1) for mono audio,
  • (L, 2) for stereo audio.

  If the sample values are floating-point numbers, they are expected to be in the [-1, 1] range; otherwise, they are assumed to be in the [-32_768, 32_767] range.

• sample_rate (int) –

  The sample rate of the audio.

• epoch (int) –

  The epoch number at which the audio is logged.

Source code in minnt/loggers/wandb_logger.py

def log_audio(self, label: str, audio: AnyArray, sample_rate: int, epoch: int) -> Self:
    audio = self.preprocess_audio(audio).numpy()
    self.run.log({label: self.wandb.Audio(audio, sample_rate=sample_rate)}, step=epoch)
    return self

log_config

log_config(config: dict[str, Any], epoch: int) -> Self

Log the given configuration dictionary at the given epoch.

Parameters:

• config (dict[str, Any]) –

  A JSON-serializable dictionary representing the configuration to log.

• epoch (int) –

  The epoch number at which the configuration is logged.

Source code in minnt/loggers/wandb_logger.py

def log_config(self, config: dict[str, Any], epoch: int) -> Self:
    self.run.config.update(dict(sorted(config.items())))
    config_json = self.format_config_as_json(config)
    self.run.log({"config": config_json} | self._maybe_as_html("config", config_json), step=epoch)
    return self

log_figure

log_figure(
    label: str,
    figure: Any,
    epoch: int,
    tight_layout: bool = True,
    close: bool = True,
) -> Self

Log the given matplotlib Figure with the given label at the given epoch.

Parameters:

• label (str) –

  The label of the logged image.

• figure (Any) –

  A matplotlib Figure.

• epoch (int) –

  The epoch number at which the image is logged.

• tight_layout (bool, default: True ) –

  Whether to apply tight layout to the figure before logging it.

• close (bool, default: True ) –

  Whether to close the figure after logging it.

Source code in minnt/loggers/wandb_logger.py

def log_figure(self, label: str, figure: Any, epoch: int, tight_layout: bool = True, close: bool = True) -> Self:
    return super().log_figure(label, figure, epoch, tight_layout, close)

log_graph

log_graph(graph: Module, data: TensorOrTensors, epoch: int) -> Self

Log the given computation graph by tracing it with the given data.

Alternatively, loggers may choose to log the graph using TorchScript, run it on the given data, or use any other mechanism they see fit.

Parameters:

• graph (Module) –

  The computation graph to log, represented as a PyTorch module.

• data (TensorOrTensors) –

  The input data to use for tracing the computation graph.

• epoch (int) –

  The epoch number at which the computation graph is logged.

Source code in minnt/loggers/wandb_logger.py

def log_graph(self, graph: torch.nn.Module, data: TensorOrTensors, epoch: int) -> Self:
    # The logging in WandB has a lot of limitations. One is that all
    # children of the graph must be executed before the graph is logged;
    # however, metrics are usually not executed during the forward pass.
    # We therefore try to exclude metrics from the logged graph.
    watched_children = torch.nn.Module()
    for name, child in graph.named_children():
        if not any(isinstance(module, Metric) for module in child.modules()):
            watched_children.add_module(name, child)
    self.run.watch(watched_children, log=None, log_graph=True)
    with self.graph_in_eval_mode(graph), torch.no_grad():
        graph(*data) if isinstance(data, (list, tuple)) else graph(data)  # Run the graph to log it.
    self.run.unwatch(watched_children)
    return self

log_image

log_image(
    label: str, image: AnyArray, epoch: int, data_format: DataFormat = "HWC"
) -> Self

Log the given image with the given label at the given epoch.

Parameters:

• label (str) –

  The label of the logged image.

• image (AnyArray) –

  The image to log, represented as a PIL image or as an array of any of the following shapes (assuming "HWC" data format):

  • (H, W) or (H, W, 1) for grayscale images,
  • (H, W, 2) for grayscale images with alpha channel,
  • (H, W, 3) for RGB images,
  • (H, W, 4) for RGBA images.

  If the pixel values are floating-point numbers, they are expected to be in the [0, 1] range; otherwise, they are assumed to be in the [0, 255] range.

• epoch (int) –

  The epoch number at which the image is logged.

• data_format (DataFormat, default: 'HWC' ) –

  The data format of the image specifying whether the channels are stored in the last dimension ("HWC", the default) or in the first dimension ("CHW"); ignored for a PIL image.

Source code in minnt/loggers/wandb_logger.py

def log_image(self, label: str, image: AnyArray, epoch: int, data_format: DataFormat = "HWC") -> Self:
    image = self.preprocess_image(image, data_format).numpy()
    self.run.log({label: self.wandb.Image(image)}, step=epoch)
    return self

log_metrics

log_metrics(
    logs: dict[str, float], epoch: int, description: str | None = None
) -> Self

Log metrics collected during a given epoch, with an optional description.

Parameters:

• logs (dict[str, float]) –

  A dictionary of logged metrics for the epoch.

• epoch (int) –

  The epoch number at which the logs were collected.

• description (str | None, default: None ) –

  An optional description of the logged metrics (used only by some loggers).

Source code in minnt/loggers/wandb_logger.py

def log_metrics(self, logs: dict[str, float], epoch: int, description: str | None = None) -> Self:
    self.run.log(logs, step=epoch)
    return self

log_text

log_text(label: str, text: str, epoch: int) -> Self

Log the given text with the given label at the given epoch.

Parameters:

• label (str) –

  The label of the logged text.

• text (str) –

  The text to log.

• epoch (int) –

  The epoch number at which the text is logged.

Source code in minnt/loggers/wandb_logger.py

def log_text(self, label: str, text: str, epoch: int) -> Self:
    self.run.log({label: text} | self._maybe_as_html(label, text), step=epoch)
    return self