| 76 | |
| 77 | |
| 78 | class TbTracker: |
| 79 | |
| 80 | name = 'tensorboard' |
| 81 | requires_logging_directory = True |
| 82 | main_process_only = True |
| 83 | |
| 84 | @on_main_process |
| 85 | def __init__(self, logging_dir, **kwargs): |
| 86 | super().__init__() |
| 87 | self.logging_dir = logging_dir |
| 88 | self.writer = SummaryWriter(self.logging_dir, **kwargs) |
| 89 | |
| 90 | @property |
| 91 | def tracker(self): |
| 92 | return self.writer |
| 93 | |
| 94 | @on_main_process |
| 95 | def add_scalar(self, tag, scalar_value, **kwargs): |
| 96 | self.writer.add_scalar(tag=tag, scalar_value=scalar_value, **kwargs) |
| 97 | |
| 98 | @on_main_process |
| 99 | def add_text(self, tag, text_string, **kwargs): |
| 100 | self.writer.add_text(tag=tag, text_string=text_string, **kwargs) |
| 101 | |
| 102 | @on_main_process |
| 103 | def add_figure(self, tag, figure, **kwargs): |
| 104 | self.writer.add_figure(tag=tag, figure=figure, **kwargs) |