Migrated project

colossalai/utils/timer.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import time
+
+from .cuda import synchronize
+
+
+class Timer:
+    '''
+    A timer object which helps to log the execution times, and provides different tools to assess the times.
+    '''
+
+    def __init__(self):
+        self._started = False
+        self._start_time = time.time()
+        self._elapsed = 0
+        self._history = []
+
+    @property
+    def has_history(self):
+        return len(self._history) != 0
+
+    def start(self):
+        '''Fisrtly synchronize cuda, reset the clock and then start the timer.
+        '''
+        self._elapsed = 0
+        synchronize()
+        self._start_time = time.time()
+        self._started = True
+
+    def stop(self, keep_in_history: bool = False):
+        '''Stop the timer and record the start-stop time interval. 
+
+        :param keep_in_history: whether does it record into history each start-stop interval, defaults to False
+        :type keep_in_history: bool, optional
+        :return: start-stop interval
+        :rtype: int
+        '''
+        synchronize()
+        end_time = time.time()
+        elapsed = end_time - self._start_time
+        if keep_in_history:
+            self._history.append(elapsed)
+        self._elapsed = elapsed
+        self._started = False
+        return elapsed
+
+    def get_history_mean(self):
+        '''mean of all history start-stop time intervals.
+
+        :return: mean of time intervals
+        :rtype: int
+        '''
+        return sum(self._history) / len(self._history)
+
+    def get_history_sum(self):
+        '''add up all the start-stop time intervals.
+
+        :return: sum of time intervals
+        :rtype: int
+        '''
+        return sum(self._history)
+
+    def get_elapsed_time(self):
+        '''return the last start-stop time interval. *use it only when timer is not in progress*
+
+        :return: the last time interval
+        :rtype: int
+        '''
+        assert not self._started, 'Timer is still in progress'
+        return self._elapsed
+
+    def reset(self):
+        '''clear up the timer and its history
+        '''
+        self._history = []
+        self._started = False
+        self._elapsed = 0
+
+
+class MultiTimer:
+    '''An object contains multiple timers
+    '''
+
+    def __init__(self, on: bool = True):
+        self._on = on
+        self._timers = dict()
+
+    def start(self, name: str):
+        '''Start namely one of the timers
+
+        :param name: timer's key
+        :type name: str
+        '''
+        if self._on:
+            if name not in self._timers:
+                self._timers[name] = Timer()
+            return self._timers[name].start()
+
+    def stop(self, name: str, keep_in_history: bool):
+        '''Stop namely one of the timers.
+
+        :param name: timer's key
+        :param keep_in_history: whether does it record into history each start-stop interval
+        :type keep_in_history: bool
+        '''
+        if self._on:
+            return self._timers[name].stop(keep_in_history)
+        else:
+            return None
+
+    def get_timer(self, name):
+        '''Get timer by its name (from multitimer)
+
+        :param name: timer's key
+        :return: timer with the name you give correctly
+        :rtype: Timer 
+        '''
+        return self._timers[name]
+
+    def reset(self, name=None):
+        '''Reset timers.
+
+        :param name: if name is designated, the named timer will be reset and others will not, defaults to None
+        '''
+        if self._on:
+            if name is not None:
+                self._timers[name].reset()
+            else:
+                for timer in self._timers:
+                    timer.reset()
+
+    def is_on(self):
+
+        return self._on
+
+    def set_status(self, mode: bool):
+        self._on = mode
+
+    def __iter__(self):
+        for name, timer in self._timers.items():
+            yield name, timer