Custom Writer for training record. Parameters: ----------- log_dir : pathlib.Path or str, path to save logs. enabled : bool, whether to enable tensorboard writer.
| 7 | import os |
| 8 | |
| 9 | class CustomWriter(object): |
| 10 | ''' |
| 11 | Custom Writer for training record. |
| 12 | Parameters: |
| 13 | ----------- |
| 14 | log_dir : pathlib.Path or str, path to save logs. |
| 15 | enabled : bool, whether to enable tensorboard writer. |
| 16 | ''' |
| 17 | def __init__(self, log_dir, enabled=True): |
| 18 | self.writer = None |
| 19 | self.selected_module = '' |
| 20 | |
| 21 | if enabled: |
| 22 | self.log_dir = str(log_dir) |
| 23 | self.stats = {} |
| 24 | if not os.path.exists(self.log_dir): |
| 25 | os.makedirs(self.log_dir, exist_ok=True) |
| 26 | |
| 27 | # Attributes to record |
| 28 | self.epoch = 0 |
| 29 | self.mode = None |
| 30 | self.timer = datetime.datetime.now() |
| 31 | self.tb_writer_funcs = { |
| 32 | 'add_scalar', 'add_scalars', |
| 33 | 'add_image', 'add_images', |
| 34 | 'add_figure', |
| 35 | 'add_audio', |
| 36 | 'add_text', |
| 37 | 'add_histogram', |
| 38 | 'add_pr_curve', |
| 39 | #'add_embedding', # TODO: problem with add_embedding |
| 40 | } |
| 41 | self.tag_mode_exceptions = {'add_histogram', 'add_embedding'} # TODO : Test these two funcs. |
| 42 | |
| 43 | def dump_stats(self): |
| 44 | with open(f"{self.log_dir}/log", "w") as f: |
| 45 | json.dump(self.stats, f, |
| 46 | indent=4, |
| 47 | ensure_ascii=False, |
| 48 | separators=(",", ": "), |
| 49 | ) |
| 50 | |
| 51 | def set_epoch(self, epoch, mode): |
| 52 | ''' |
| 53 | Execute this function to update the step attribute and compute the cost time of one epoch in seconds. |
| 54 | Recommend to run this function every step. |
| 55 | This function MUST be executed before other custom writer functions. |
| 56 | Parameters: |
| 57 | ------------ |
| 58 | step : int, step number. |
| 59 | mode : str, 'train' or 'valid' |
| 60 | ''' |
| 61 | if epoch == 0: |
| 62 | self.timer = datetime.datetime.now() |
| 63 | elif epoch != self.epoch: |
| 64 | duration = datetime.datetime.now() - self.timer |
| 65 | second_per_epoch = duration.total_seconds() / (epoch - self.epoch) |
| 66 | self.add_scalar(tag='second_per_epoch', data=second_per_epoch) |