mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
5c4545f857
Reproduce command: ```bash ghstack checkout https://github.com/pytorch/pytorch/pull/138447 git checkout HEAD~1 torch/ lintrunner -a --take "PYFMT" --all-files ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/138447 Approved by: https://github.com/ezyang
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from torch._C._monitor import * # noqa: F403
|
|
from torch._C._monitor import _WaitCounter # type: ignore[attr-defined]
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
|
|
STAT_EVENT = "torch.monitor.Stat"
|
|
|
|
|
|
class TensorboardEventHandler:
|
|
"""
|
|
TensorboardEventHandler is an event handler that will write known events to
|
|
the provided SummaryWriter.
|
|
|
|
This currently only supports ``torch.monitor.Stat`` events which are logged
|
|
as scalars.
|
|
|
|
Example:
|
|
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_MONITOR)
|
|
>>> # xdoctest: +REQUIRES(module:tensorboard)
|
|
>>> from torch.utils.tensorboard import SummaryWriter
|
|
>>> from torch.monitor import TensorboardEventHandler, register_event_handler
|
|
>>> writer = SummaryWriter("log_dir")
|
|
>>> register_event_handler(TensorboardEventHandler(writer))
|
|
"""
|
|
|
|
def __init__(self, writer: "SummaryWriter") -> None:
|
|
"""
|
|
Constructs the ``TensorboardEventHandler``.
|
|
"""
|
|
self._writer = writer
|
|
|
|
def __call__(self, event: Event) -> None:
|
|
if event.name == STAT_EVENT:
|
|
for k, v in event.data.items():
|
|
self._writer.add_scalar(k, v, walltime=event.timestamp.timestamp())
|