Summarize the average GPU utilization within an epoch. It will start a process to obtain GPU utilization through NVML every second within the epoch (the trigger_epoch time was not included), and write average utilization to monitors. This callback creates a process, therefore it's
| 23 | |
| 24 | |
| 25 | class GPUUtilizationTracker(Callback): |
| 26 | """ Summarize the average GPU utilization within an epoch. |
| 27 | |
| 28 | It will start a process to obtain GPU utilization through NVML every second |
| 29 | within the epoch (the trigger_epoch time was not included), |
| 30 | and write average utilization to monitors. |
| 31 | |
| 32 | This callback creates a process, therefore it's not safe to be used with MPI. |
| 33 | """ |
| 34 | |
| 35 | _chief_only = False |
| 36 | |
| 37 | def __init__(self, devices=None): |
| 38 | """ |
| 39 | Args: |
| 40 | devices (list[int]): physical GPU ids to monitor. If None, will guess from the environment. |
| 41 | """ |
| 42 | assert os.name != 'nt', "GPUUtilizationTracker does not support windows!" |
| 43 | self._devices = devices |
| 44 | self._enabled = True |
| 45 | |
| 46 | def _guess_devices(self): |
| 47 | env = os.environ.get('CUDA_VISIBLE_DEVICES') |
| 48 | if env is None: |
| 49 | devices = list(range(get_num_gpu())) |
| 50 | if len(devices) > 1: |
| 51 | logger.warn("[GPUUtilizationTracker] Both devices and CUDA_VISIBLE_DEVICES are None! " |
| 52 | "Will monitor all {} visible GPUs!".format(len(devices))) |
| 53 | else: |
| 54 | if len(env): |
| 55 | devices = list(map(int, env.split(','))) |
| 56 | else: |
| 57 | devices = [] |
| 58 | return devices |
| 59 | |
| 60 | def _setup_graph(self): |
| 61 | # special heuristics for Horovod |
| 62 | from ..train import HorovodTrainer |
| 63 | if isinstance(self.trainer, HorovodTrainer): |
| 64 | if self.trainer.mpi_enabled(): |
| 65 | logger.warn("GPUUtilizationTracker is disabled under MPI.") |
| 66 | self._enabled = False |
| 67 | return |
| 68 | else: |
| 69 | self._devices = [self.trainer.hvd.local_rank()] |
| 70 | |
| 71 | if self._devices is None: |
| 72 | self._devices = self._guess_devices() |
| 73 | assert len(self._devices), "[GPUUtilizationTracker] No GPU device given!" |
| 74 | |
| 75 | self._evt = mp.Event() |
| 76 | self._stop_evt = mp.Event() |
| 77 | self._queue = mp.Queue() |
| 78 | self._proc = mp.Process(target=self.worker, args=( |
| 79 | self._evt, self._queue, self._stop_evt, self._devices)) |
| 80 | ensure_proc_terminate(self._proc) |
| 81 | start_proc_mask_signal(self._proc) |
| 82 |
no outgoing calls
no test coverage detected