diff --git a/src/proloaf/event_logging.py b/src/proloaf/event_logging.py index 0d2962245e7727364d1dd701bac03bf6c126ae7b..567816232f24116532278017350ea584f0233dd8 100644 --- a/src/proloaf/event_logging.py +++ b/src/proloaf/event_logging.py @@ -1,10 +1,13 @@ import os import logging import logging.config +import functools +import time from proloaf.confighandler import read_config MAIN_PATH = os.path.dirname(os.path.abspath(__file__)) + def create_event_logger( name: str, config_path: os.path = os.path.join(MAIN_PATH, 'event_logging_conf.json'), @@ -56,3 +59,22 @@ def create_event_logger( logger = logging.getLogger(name) return logger + + +def timer(logger: logging.Logger): + """ + Decorator which passes the runtime of the decorated function to a logger. + The logger then prints the name and runtime of the function, if the logger is specified to print statements + of rank "INFO". + """ + def decorator(func): + @functools.wraps(func) + def wrapper_timer(*args, **kwargs): + start_time = time.perf_counter() + value = func(*args, **kwargs) + end_time = time.perf_counter() + run_time = end_time - start_time + logger.info(f"Finished {func.__name__!r} in {run_time:.4f} secs") + return value + return wrapper_timer + return decorator diff --git a/src/proloaf/modelhandler.py b/src/proloaf/modelhandler.py index 5b4b77c9ba1dc2b42958e206f4f087483d1e4bbc..eb9e82f09f0aa27ec1b01e94fc9416a2ffe19e6f 100644 --- a/src/proloaf/modelhandler.py +++ b/src/proloaf/modelhandler.py @@ -45,6 +45,7 @@ from proloaf.loghandler import ( end_tensorboard, ) from proloaf.event_logging import create_event_logger +from proloaf.event_logging import timer logger = create_event_logger(__name__) @@ -861,6 +862,7 @@ class ModelHandler: end_tensorboard(tb, hparams, values) return temp_model_wrap + @timer(logger) def fit( self, train_data: proloaf.tensorloader.TimeSeriesData, @@ -900,6 +902,7 @@ class ModelHandler: ) return self + @timer(logger) def predict(self, inputs_enc: torch.Tensor, inputs_dec: torch.Tensor): """ Get the predictions for the given model and data