/
timer.py
63 lines (49 loc) · 1.88 KB
/
timer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import time
import numpy
from chainer import cuda
from chainer import function
class TimerHook(function.FunctionHook):
"""Function hook for measuring elapsed time of functions.
Attributes:
call_history: List of measurement results. It consists of pairs of
the function that calls this hook and the elapsed time
the function consumes.
"""
name = 'TimerHook'
def __init__(self):
self.call_history = []
def _preprocess(self):
if self.xp == numpy:
self.start = time.time()
else:
self.start = cuda.Event()
self.stop = cuda.Event()
self.start.record()
def forward_preprocess(self, function, in_data):
self.xp = cuda.get_array_module(*in_data)
self._preprocess()
def backward_preprocess(self, function, in_data, out_grad):
self.xp = cuda.get_array_module(*(in_data + out_grad))
self._preprocess()
def _postprocess(self, function):
if self.xp == numpy:
self.stop = time.time()
elapsed_time = self.stop - self.start
else:
self.stop.record()
self.stop.synchronize()
# Note that `get_elapsed_time` returns result in milliseconds
elapsed_time = cuda.cupy.cuda.get_elapsed_time(
self.start, self.stop) / 1000
self.call_history.append((function, elapsed_time))
def forward_postprocess(self, function, in_data):
xp = cuda.get_array_module(*in_data)
assert xp == self.xp
self._postprocess(function)
def backward_postprocess(self, function, in_data, out_grad):
xp = cuda.get_array_module(*(in_data + out_grad))
assert xp == self.xp
self._postprocess(function)
def total_time(self):
"""Returns total elapsed time in seconds."""
return sum(t for (_, t) in self.call_history)