pytorch

1
from torch._C._monitor import *  # noqa: F403
2
from typing import TYPE_CHECKING
3

4
from torch._C._monitor import _WaitCounter  # type: ignore[attr-defined]
5

6
if TYPE_CHECKING:
7
    from torch.utils.tensorboard import SummaryWriter
8

9

10
STAT_EVENT = "torch.monitor.Stat"
11

12

13
class TensorboardEventHandler:
14
    """
15
    TensorboardEventHandler is an event handler that will write known events to
16
    the provided SummaryWriter.
17

18
    This currently only supports ``torch.monitor.Stat`` events which are logged
19
    as scalars.
20

21
    Example:
22
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_MONITOR)
23
        >>> # xdoctest: +REQUIRES(module:tensorboard)
24
        >>> from torch.utils.tensorboard import SummaryWriter
25
        >>> from torch.monitor import TensorboardEventHandler, register_event_handler
26
        >>> writer = SummaryWriter("log_dir")
27
        >>> register_event_handler(TensorboardEventHandler(writer))
28
    """
29
    def __init__(self, writer: "SummaryWriter") -> None:
30
        """
31
        Constructs the ``TensorboardEventHandler``.
32
        """
33
        self._writer = writer
34

35
    def __call__(self, event: Event) -> None:
36
        if event.name == STAT_EVENT:
37
            for k, v in event.data.items():
38
                self._writer.add_scalar(k, v, walltime=event.timestamp.timestamp())
39