| 3 | |
| 4 | |
| 5 | class WriterTensorboardX(): |
| 6 | def __init__(self, writer_dir, logger, enable): |
| 7 | self.writer = None |
| 8 | if enable: |
| 9 | log_path = writer_dir |
| 10 | try: |
| 11 | self.writer = importlib.import_module('tensorboardX').SummaryWriter(log_path) |
| 12 | except ModuleNotFoundError: |
| 13 | message = """ |
| 14 | TensorboardX visualization is configured to use, but currently not installed on this machine. |
| 15 | Please install the package by 'pip install tensorboardx' command or turn off the option |
| 16 | in the 'config.json' file. |
| 17 | """ |
| 18 | warnings.warn(message, UserWarning) |
| 19 | # logger.warn() |
| 20 | self.step = 0 |
| 21 | self.mode = '' |
| 22 | |
| 23 | self.tensorboard_writer_ftns = [ |
| 24 | 'add_scalar', 'add_scalars', 'add_image', 'add_audio', 'add_text', 'add_histogram', |
| 25 | 'add_pr_curve', 'add_embedding' |
| 26 | ] |
| 27 | |
| 28 | def set_step(self, step, mode='train'): |
| 29 | self.mode = mode |
| 30 | self.step = step |
| 31 | |
| 32 | def __getattr__(self, name): |
| 33 | """ |
| 34 | If visualization is configured to use: |
| 35 | return add_data() methods of tensorboard with additional information (step, tag) added. |
| 36 | Otherwise: |
| 37 | return blank function handle that does nothing |
| 38 | """ |
| 39 | if name in self.tensorboard_writer_ftns: |
| 40 | add_data = getattr(self.writer, name, None) |
| 41 | |
| 42 | def wrapper(tag, data, *args, **kwargs): |
| 43 | if add_data is not None: |
| 44 | add_data('{}/{}'.format(self.mode, tag), data, self.step, *args, **kwargs) |
| 45 | return wrapper |
| 46 | else: |
| 47 | # default action for returning methods defined in this class, set_step() for instance. |
| 48 | try: |
| 49 | attr = object.__getattr__(name) |
| 50 | except AttributeError: |
| 51 | raise AttributeError("type object 'WriterTensorboardX' has no attribute '{}'".format(name)) |
| 52 | return attr |