import time
import numpy
from chainer import cuda
from chainer import function
[docs]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)
[docs] def total_time(self):
"""Returns total elapsed time in seconds."""
return sum(t for (_, t) in self.call_history)