Construc tensorboard writer (self.writer). The tensorboard is saved at os.path.join(tb_dir, file_name).
| 131 | |
| 132 | |
| 133 | class TBLog: |
| 134 | """ |
| 135 | Construc tensorboard writer (self.writer). |
| 136 | The tensorboard is saved at os.path.join(tb_dir, file_name). |
| 137 | """ |
| 138 | |
| 139 | def __init__(self, tb_dir, file_name, use_tensorboard=False): |
| 140 | self.tb_dir = tb_dir |
| 141 | self.use_tensorboard = use_tensorboard |
| 142 | if self.use_tensorboard: |
| 143 | self.writer = SummaryWriter(os.path.join(self.tb_dir, file_name)) |
| 144 | else: |
| 145 | self.writer = CustomWriter(os.path.join(self.tb_dir, file_name)) |
| 146 | |
| 147 | def update(self, tb_dict, it, suffix=None, mode="train"): |
| 148 | """ |
| 149 | Args |
| 150 | tb_dict: contains scalar values for updating tensorboard |
| 151 | it: contains information of iteration (int). |
| 152 | suffix: If not None, the update key has the suffix. |
| 153 | """ |
| 154 | if suffix is None: |
| 155 | suffix = '' |
| 156 | if self.use_tensorboard: |
| 157 | for key, value in tb_dict.items(): |
| 158 | self.writer.add_scalar(suffix + key, value, it) |
| 159 | else: |
| 160 | self.writer.set_epoch(it, mode) |
| 161 | for key, value in tb_dict.items(): |
| 162 | self.writer.add_scalar(suffix + key, value) |
| 163 | self.writer.plot_stats() |
| 164 | self.writer.dump_stats() |
| 165 | |
| 166 | |
| 167 | class AverageMeter(object): |
no outgoing calls
no test coverage detected