|
from collections import deque, defaultdict |
|
from functools import wraps |
|
from types import GeneratorType |
|
from typing import Callable |
|
import numpy as np |
|
import time |
|
from ditk import logging |
|
from ding.framework import task |
|
|
|
|
|
class StepTimer: |
|
|
|
def __init__(self, print_per_step: int = 1, smooth_window: int = 10) -> None: |
|
""" |
|
Overview: |
|
Print time cost of each step (execute one middleware). |
|
Arguments: |
|
- print_per_step (:obj:`int`): Print each N step. |
|
- smooth_window (:obj:`int`): The window size to smooth the mean. |
|
""" |
|
|
|
self.print_per_step = print_per_step |
|
self.records = defaultdict(lambda: deque(maxlen=print_per_step * smooth_window)) |
|
|
|
def __call__(self, fn: Callable) -> Callable: |
|
step_name = getattr(fn, "__name__", type(fn).__name__) |
|
|
|
@wraps(fn) |
|
def executor(ctx): |
|
start_time = time.time() |
|
time_cost = 0 |
|
g = fn(ctx) |
|
if isinstance(g, GeneratorType): |
|
try: |
|
next(g) |
|
except StopIteration: |
|
pass |
|
time_cost = time.time() - start_time |
|
yield |
|
start_time = time.time() |
|
try: |
|
next(g) |
|
except StopIteration: |
|
pass |
|
time_cost += time.time() - start_time |
|
else: |
|
time_cost = time.time() - start_time |
|
self.records[step_name].append(time_cost) |
|
if ctx.total_step % self.print_per_step == 0: |
|
logging.info( |
|
"[Step Timer][Node:{:>2}] {}: Cost: {:.2f}ms, Mean: {:.2f}ms".format( |
|
task.router.node_id or 0, step_name, time_cost * 1000, |
|
np.mean(self.records[step_name]) * 1000 |
|
) |
|
) |
|
|
|
return executor |
|
|
