| 5 | |
| 6 | |
| 7 | class Logger(object): |
| 8 | |
| 9 | def __init__(self, log_dir=None, accelerator=None) -> None: |
| 10 | self.log_dir = log_dir |
| 11 | self.accelerator = accelerator |
| 12 | |
| 13 | if self.log_dir is not None: |
| 14 | self.txt_writer = open(os.path.join(self.log_dir, 'logger.log'), 'a') |
| 15 | else: |
| 16 | self.txt_writer = None |
| 17 | |
| 18 | if SummaryWriter is not None and self.accelerator.is_main_process: |
| 19 | self.writer = SummaryWriter(self.log_dir) |
| 20 | else: |
| 21 | self.writer = None |
| 22 | |
| 23 | def log_scalars(self, scalar_dict, step, prefix=None): |
| 24 | if self.writer is None: |
| 25 | return |
| 26 | for k in scalar_dict: |
| 27 | v = scalar_dict[k] |
| 28 | if isinstance(v, torch.Tensor): |
| 29 | v = v.detach().cpu().item() |
| 30 | if prefix is not None: |
| 31 | k = prefix + '_' + k |
| 32 | self.writer.add_scalar(k, v, step) |
| 33 | return |
| 34 | |
| 35 | def log_messages(self, message: str): |
| 36 | if self.txt_writer is not None: |
| 37 | self.txt_writer.write(message + "\n") |
| 38 | self.txt_writer.flush() |
| 39 | print(message, flush=True) |
| 40 | return |
| 41 | |