简单日志分析器

这是一个用作训练器应用程序示例一部分的简单分析器。它将 Lightning 训练阶段持续时间记录到 Tensorboard 等日志记录器中。此输出用于 Ax 的 HPO 优化。

import time
from typing import Dict

from pytorch_lightning.loggers.logger import Logger

from pytorch_lightning.profilers.profiler import Profiler


class SimpleLoggingProfiler(Profiler):
    """
    This profiler records the duration of actions (in seconds) and reports the
    mean duration of each action to the specified logger. Reported metrics are
    in the format `duration_<event>`.
    """

    def __init__(self, logger: Logger) -> None:
        super().__init__()

        self.current_actions: Dict[str, float] = {}
        self.logger = logger

    def start(self, action_name: str) -> None:
        if action_name in self.current_actions:
            raise ValueError(
                f"Attempted to start {action_name} which has already started."
            )
        self.current_actions[action_name] = time.monotonic()

    def stop(self, action_name: str) -> None:
        end_time = time.monotonic()
        if action_name not in self.current_actions:
            raise ValueError(
                f"Attempting to stop recording an action ({action_name}) which was never started."
            )
        start_time = self.current_actions.pop(action_name)
        duration = end_time - start_time
        self.logger.log_metrics({"duration_" + action_name: duration})

    def summary(self) -> str:
        return ""


# sphinx_gallery_thumbnail_path = '_static/img/gallery-lib.png'