setup logger for training and testing. Args: save_dir(str): location to save log file distributed_rank(int): device rank when multi-gpu environment filename (string): log save name. mode(str): log file write mode, `append` or `override`. default is `a`. Retur
(save_dir, distributed_rank=0, filename="log.txt", mode="a")
| 62 | |
| 63 | |
| 64 | def setup_logger(save_dir, distributed_rank=0, filename="log.txt", mode="a"): |
| 65 | """setup logger for training and testing. |
| 66 | Args: |
| 67 | save_dir(str): location to save log file |
| 68 | distributed_rank(int): device rank when multi-gpu environment |
| 69 | filename (string): log save name. |
| 70 | mode(str): log file write mode, `append` or `override`. default is `a`. |
| 71 | |
| 72 | Return: |
| 73 | logger instance. |
| 74 | """ |
| 75 | loguru_format = ( |
| 76 | "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | " |
| 77 | "<level>{level: <8}</level> | " |
| 78 | "<cyan>{name}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>" |
| 79 | ) |
| 80 | |
| 81 | logger.remove() |
| 82 | save_file = os.path.join(save_dir, filename) |
| 83 | if mode == "o" and os.path.exists(save_file): |
| 84 | os.remove(save_file) |
| 85 | # only keep logger in rank0 process |
| 86 | if distributed_rank == 0: |
| 87 | logger.add( |
| 88 | sys.stderr, |
| 89 | format=loguru_format, |
| 90 | level="INFO", |
| 91 | enqueue=True, |
| 92 | ) |
| 93 | logger.add(save_file) |
| 94 | |
| 95 | # redirect stdout/stderr to loguru |
| 96 | redirect_sys_output("INFO") |