Method worker

tensorpack/callbacks/prof.py:122–159

Args: devices (list[int])

(evt, rst_queue, stop_evt, devices)

    @staticmethod
    def worker(evt, rst_queue, stop_evt, devices):
        """
        Args:
            devices (list[int])
        """
        try:
            with NVMLContext() as ctx:
                devices = [ctx.device(i) for i in devices]
                while True:
                    evt.wait()  # start epoch
                    evt.clear()
                    if stop_evt.is_set():   # or on exit
                        return

                    stats = np.zeros((len(devices),), dtype='f4')
                    cnt = 0
                    while True:
                        time.sleep(1)

                        data = [d.utilization()['gpu'] for d in devices]
                        data = list(map(float, data))
                        stats += data
                        cnt += 1

                        if evt.is_set():    # stop epoch
                            if stop_evt.is_set():   # or on exit
                                return
                            if cnt > 1:
                                # Ignore the last datapoint. Usually is zero, makes us underestimate the util.
                                stats -= data
                                cnt -= 1
                            rst_queue.put(stats / cnt)
                            evt.clear()
                            break
        except Exception:
            logger.exception("Exception in GPUUtilizationTracker.worker")
            rst_queue.put(-1)
            return


# Can add more features from tfprof