| |
| |
|
|
| |
| """ |
| Support different forms of monitoring such as wandb and tensorboard |
| """ |
|
|
| from abc import ABC, abstractmethod |
| import deepspeed.comm as dist |
|
|
|
|
| class Monitor(ABC): |
|
|
| @abstractmethod |
| def __init__(self, monitor_config): |
| self.monitor_config = monitor_config |
|
|
| @abstractmethod |
| def write_events(self, event_list): |
| pass |
|
|
|
|
| from .wandb import WandbMonitor |
| from .tensorboard import TensorBoardMonitor |
| from .csv_monitor import csvMonitor |
| from .comet import CometMonitor |
|
|
|
|
| class MonitorMaster(Monitor): |
|
|
| def __init__(self, monitor_config): |
| super().__init__(monitor_config) |
| self.tb_monitor = None |
| self.wandb_monitor = None |
| self.csv_monitor = None |
| self.comet_monitor = None |
| self.enabled = monitor_config.enabled |
|
|
| if dist.get_rank() == 0: |
| if monitor_config.tensorboard.enabled: |
| self.tb_monitor = TensorBoardMonitor(monitor_config.tensorboard) |
| if monitor_config.wandb.enabled: |
| self.wandb_monitor = WandbMonitor(monitor_config.wandb) |
| if monitor_config.csv_monitor.enabled: |
| self.csv_monitor = csvMonitor(monitor_config.csv_monitor) |
| if monitor_config.comet.enabled: |
| self.comet_monitor = CometMonitor(monitor_config.comet) |
|
|
| def write_events(self, event_list): |
| if dist.get_rank() == 0: |
| if self.tb_monitor is not None: |
| self.tb_monitor.write_events(event_list) |
| if self.wandb_monitor is not None: |
| self.wandb_monitor.write_events(event_list) |
| if self.csv_monitor is not None: |
| self.csv_monitor.write_events(event_list) |
| if self.comet_monitor is not None: |
| self.comet_monitor.write_events(event_list) |
|
|
