File size: 1,071 Bytes
cd47a59 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | from torch.utils.tensorboard import SummaryWriter
from accelerate.tracking import GeneralTracker, on_main_process
import os
from typing import Union
class TbTracker(GeneralTracker):
name = "tensorboard"
requires_logging_directory = True
@on_main_process
def __init__(self, run_name: str, logging_dir: Union[str, os.PathLike],
**kwargs):
super().__init__()
self.run_name = run_name
self.logging_dir = os.path.join(logging_dir, run_name)
self.writer = SummaryWriter(self.logging_dir, **kwargs)
@property
def tracker(self):
return self.writer
@on_main_process
def add_scalar(self, tag, scalar_value, **kwargs):
self.writer.add_scalar(tag=tag, scalar_value=scalar_value, **kwargs)
@on_main_process
def add_text(self, tag, text_string, **kwargs):
self.writer.add_text(tag=tag, text_string=text_string, **kwargs)
@on_main_process
def add_figure(self, tag, figure, **kwargs):
self.writer.add_figure(tag=tag, figure=figure, **kwargs)
|