(self, save_dir: str, log_level: str, rank: int = 0, mode='train')
| 43 | OmegaConf.save(config, out_path) |
| 44 | |
| 45 | def setup_save(self, save_dir: str, log_level: str, rank: int = 0, mode='train'): |
| 46 | log_level_str = log_level.upper() |
| 47 | self.save_dir = save_dir |
| 48 | self.ckpt_dir = osp.join(save_dir, 'checkpointing') |
| 49 | os.makedirs(self.ckpt_dir, exist_ok=True) |
| 50 | # self.model_weight_dir = osp.join(save_dir, 'weights') |
| 51 | # os.makedirs(self.model_weight_dir, exist_ok=True) |
| 52 | if mode == 'train': |
| 53 | self.vis_dir = osp.join(save_dir, 'visualization') |
| 54 | else: |
| 55 | self.vis_dir = osp.join(save_dir, 'test_visualization') |
| 56 | os.makedirs(self.vis_dir, exist_ok=True) |
| 57 | self.logfile = osp.join(save_dir, 'train.log') |
| 58 | if rank == 0: |
| 59 | handler = logging.FileHandler(self.logfile, mode='a') |
| 60 | handler.setLevel(log_level_str) |
| 61 | if dist.is_initialized(): |
| 62 | dist.barrier() |
| 63 | if rank != 0: |
| 64 | handler = logging.FileHandler(self.logfile, mode='a') |
| 65 | handler.setLevel(log_level_str) |
| 66 | formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s') |
| 67 | handler.setFormatter(formatter) |
| 68 | console_handler = logging.StreamHandler() |
| 69 | console_handler.setLevel(log_level_str) |
| 70 | console_handler.setFormatter(formatter) |
| 71 | logger = get_logger(__name__, log_level=log_level_str) |
| 72 | logger.logger.addHandler(handler) |
| 73 | logger.logger.addHandler(console_handler) |
| 74 | self.logger = logger |
| 75 | try: |
| 76 | self.tb_tracker = TbTracker(save_dir) |
| 77 | except ValueError: |
| 78 | pass |
| 79 | |
| 80 | # @staticmethod |
| 81 | # def load_ckpt(model, optimizer=None, ckpt_dir:str=None): |
no test coverage detected