pytorch

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

3
from typing import TYPE_CHECKING
4

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

8

9
STAT_EVENT = "torch.monitor.Stat"
10

11

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

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

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

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