Added Aim logging, integration and Demo

[osoblanco]

Pull Request FAQ

• documentation
• contribution guide
• minimal examples section
• changelog for main framework updates
• Catalyst slack (#__questions channel) for issue discussion

Description

Added an Aim integration (completed pipeline, docs updated) with a demo within the examples.

Related Issue

Type of Change

• [x] Examples / docs / tutorials / contributors update
• [ ] Bug fix (non-breaking change which fixes an issue)
• [x] Improvement (non-breaking change which improves an existing feature)
• [x] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to change)

PR review

Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Checklist

• [x] Have you updated tests for the new functionality?
• [x] Have you added your new classes/functions to the docs?
• [ ] Have you updated the CHANGELOG?
• [x] Have you run colab minimal CI/CD with latest requirements? Please attach the notebook link.
• [x] Have you run colab minimal CI/CD with minimal requirements? Please attach the notebook link.

[pep8speaks]

Hello @osoblanco! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers:

Comment last updated at 2022-08-17 10:02:55 UTC

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

:white_check_mark: osoblanco
:x: tmynn
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Scitator]

@osoblanco could you please check the codestyle?

[tamohannes]

hey @Scitator, sorry for a late response, just resolved the codestyle issues. Could you please take a look at it and run the workflows?

[YodaEmbedding]

I'm also an Aim user. Below is what I have been using in my own projects, in case it helps in taking the best of the two approaches. (EDIT: Looks like this PR may have been based off some older code that I published a few months ago.)

A couple of things to note:

• Consider running black -l 79 to format more consistently. Remove inconsistent whitespace, unused print comments, etc.
• Consider using the standard conventions. For instance, use "loader" and "scope":

    if loader_key is not None:
        context["loader"] = loader_key
    if scope is not None:
        context["scope"] = scope

• Add support for log_figure.

---

My latest approach:

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

import aim
import numpy as np
from catalyst.core.logger import ILogger
from catalyst.settings import SETTINGS

if TYPE_CHECKING:
    from catalyst.core.runner import IRunner


class AimLogger(ILogger):
    """Aim logger for parameters, metrics, images and other artifacts.

    Aim documentation: https://aimstack.readthedocs.io/en/latest/.

    Args:
        experiment: Name of the experiment in Aim to log to.
        run_hash: Run hash.
        exclude: Name of key to exclude from logging.
        log_batch_metrics: boolean flag to log batch metrics
            (default: SETTINGS.log_batch_metrics or False).
        log_epoch_metrics: boolean flag to log epoch metrics
            (default: SETTINGS.log_epoch_metrics or True).

    Python API examples:

    .. code-block:: python

        from catalyst import dl

        runner = dl.SupervisedRunner()
        runner.train(
            ...,
            loggers={"aim": dl.AimLogger(experiment="test_exp")}
        )

    .. code-block:: python

        from catalyst import dl

        class CustomRunner(dl.IRunner):
            # ...

            def get_loggers(self):
                return {
                    "console": dl.ConsoleLogger(),
                    "aim": dl.AimLogger(experiment="test_exp")
                }

            # ...

        runner = CustomRunner().run()
    """

    exclude: List[str]
    run: aim.Run

    def __init__(
        self,
        *,
        experiment: Optional[str] = None,
        run_hash: Optional[str] = None,
        exclude: Optional[List[str]] = None,
        log_batch_metrics: bool = SETTINGS.log_batch_metrics,
        log_epoch_metrics: bool = SETTINGS.log_epoch_metrics,
        repo: Optional[Union[str, aim.Repo]] = None,
        **kwargs,
    ) -> None:
        super().__init__(
            log_batch_metrics=log_batch_metrics,
            log_epoch_metrics=log_epoch_metrics,
        )
        self.exclude = [] if exclude is None else exclude
        self.run = aim.Run(
            run_hash=run_hash,
            repo=repo,
            experiment=experiment,
            **kwargs,
        )

    @property
    def logger(self):
        """Internal logger/experiment/etc. from the monitoring system."""
        return self.run

    def log_artifact(
        self,
        tag: str,
        runner: "IRunner",
        artifact: object = None,
        path_to_artifact: Optional[str] = None,
        scope: Optional[str] = None,
        kind: str = "text",
        artifact_kwargs: Dict[str, Any] = {},
    ) -> None:
        """Logs a local file or directory as an artifact to the logger."""
        if path_to_artifact:
            mode = "r" if kind == "text" else "rb"
            with open(path_to_artifact, mode) as f:
                artifact = f.read()
        kind_dict = {
            "audio": aim.Audio,
            "figure": aim.Figure,
            "image": aim.Image,
            "text": aim.Text,
        }
        value = kind_dict[kind](artifact, **artifact_kwargs)
        context, kwargs = _aim_context(runner, scope)
        self.run.track(value, tag, context=context, **kwargs)

    def log_image(
        self,
        tag: str,
        image,
        runner: "IRunner",
        scope: Optional[str] = None,
        image_kwargs: Dict[str, Any] = {},
    ) -> None:
        """Logs image to Aim for current scope on current step."""
        value = aim.Image(image, **image_kwargs)
        context, kwargs = _aim_context(runner, scope)
        self.run.track(value, tag, context=context, **kwargs)

    def log_hparams(self, hparams: Dict, runner: "Optional[IRunner]" = None) -> None:
        """Logs parameters for current scope.

        Args:
            hparams: Parameters to log.
            runner: experiment runner
        """
        d: dict[str, Any] = {}
        _build_params_dict(hparams, d, self.exclude)
        for k, v in d.items():
            self.run[k] = v

    def log_metrics(
        self,
        metrics: Dict[str, float],
        scope: str,
        runner: "IRunner",
    ) -> None:
        """Logs batch and epoch metrics to Aim."""
        if scope == "batch" and self.log_batch_metrics:
            metrics = {k: float(v) for k, v in metrics.items()}
            self._log_metrics(
                metrics=metrics,
                runner=runner,
                loader_key=runner.loader_key,
                scope=scope,
            )
        elif scope == "epoch" and self.log_epoch_metrics:
            for loader_key, per_loader_metrics in metrics.items():
                self._log_metrics(
                    metrics=per_loader_metrics,  # type: ignore
                    runner=runner,
                    loader_key=loader_key,
                    scope=scope,
                )

    def log_figure(
        self,
        tag: str,
        fig: Any,
        runner: "IRunner",
        scope: Optional[str] = None,
        kwargs: Dict[str, Any] = {},
    ) -> None:
        """Logs figure to Aim for current scope on current step."""
        value = aim.Figure(fig, **kwargs)
        context, kwargs = _aim_context(runner, scope)
        self.run.track(value, tag, context=context, **kwargs)

    def close_log(self) -> None:
        """End an active Aim run."""
        self.run.close()

    def _log_metrics(
        self,
        metrics: Dict[str, float],
        runner: "IRunner",
        loader_key: str,
        scope: str = "",
    ):
        context, kwargs = _aim_context(runner, scope, loader_key)
        for key, value in metrics.items():
            self.run.track(value, key, context=context, **kwargs)


def _aim_context(
    runner: "IRunner",
    scope: Optional[str],
    loader_key: Optional[str] = None,
    all_scope_steps: bool = False,
):
    if loader_key is None:
        loader_key = runner.loader_key
    context = {}
    if loader_key is not None:
        context["loader"] = loader_key
    if scope is not None:
        context["scope"] = scope
    kwargs = {}
    if all_scope_steps or scope == "batch":
        kwargs["step"] = runner.batch_step
    if all_scope_steps or scope == "epoch" or scope == "loader":
        kwargs["epoch"] = runner.epoch_step
    return context, kwargs


def _build_params_dict(
    dictionary: Dict[str, Any],
    prefix: Dict[str, Any],
    exclude: List[str],
):
    for name, value in dictionary.items():
        if name in exclude:
            continue

        if isinstance(value, dict):
            if name not in prefix:
                prefix[name] = {}
            _build_params_dict(value, prefix[name], exclude)
        else:
            prefix[name] = value

[tamohannes]

Hey @YodaEmbedding thanks a lot for sharing the code snippet! I have refactored the logger and used your code suggestions. Please let me know if there is something that I missed.

[tamohannes]

Hey @Scitator could you please run the workflows once again. I have formatted the directory files using the command suggested by @YodaEmbedding: black catalyst -l 79

[YodaEmbedding]

I think that reformatted the entire repository. Try:

commit_hash="$(git rev-parse HEAD)"
git reset HEAD~
git checkout .
git checkout "${commit_hash}" -- catalyst/loggers/aim.py
git add catalyst/loggers/aim.py
git commit -m "[fix] Refactoring AimLogger. fix PEP-8 compliance errors"
git show
git push --force

[tamohannes]

Hey @YodaEmbedding sorry for late reply. That's right I accidentally formatted the entire repo. Thanks a lot for the clear steps, it did the job!

[tamohannes]

Hey @Scitator could you please run the workflows once again.

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[tamohannes]

Hey @Scitator please run the workflows, so we can act if further corrections are needed.

[tamohannes]

Hi, @Scitator this PR got closed automatically 😕 Let's reopen it, and run the checks so we can merge it 🚀 I will love to work on the issues, so we can give the users the possibility to track their experiments with the tool they love ☺️