diff --git a/monai/handlers/stats_handler.py b/monai/handlers/stats_handler.py index 58917e666b..8471a87e8e 100644 --- a/monai/handlers/stats_handler.py +++ b/monai/handlers/stats_handler.py @@ -66,8 +66,8 @@ class StatsHandler: def __init__( self, - iteration_log: bool = True, - epoch_log: bool = True, + iteration_log: bool | Callable[[Engine, int], bool] = True, + epoch_log: bool | Callable[[Engine, int], bool] = True, epoch_print_logger: Callable[[Engine], Any] | None = None, iteration_print_logger: Callable[[Engine], Any] | None = None, output_transform: Callable = lambda x: x[0], @@ -80,8 +80,14 @@ def __init__( """ Args: - iteration_log: whether to log data when iteration completed, default to `True`. - epoch_log: whether to log data when epoch completed, default to `True`. + iteration_log: whether to log data when iteration completed, default to `True`. ``iteration_log`` can + be also a function and it will be interpreted as an event filter + (see https://pytorch.org/ignite/generated/ignite.engine.events.Events.html for details). + Event filter function accepts as input engine and event value (iteration) and should return True/False. + Event filtering can be helpful to customize iteration logging frequency. + epoch_log: whether to log data when epoch completed, default to `True`. ``epoch_log`` can be + also a function and it will be interpreted as an event filter. See ``iteration_log`` argument for more + details. epoch_print_logger: customized callable printer for epoch level logging. Must accept parameter "engine", use default printer if None. iteration_print_logger: customized callable printer for iteration level logging. @@ -135,9 +141,15 @@ def attach(self, engine: Engine) -> None: " please call `logging.basicConfig(stream=sys.stdout, level=logging.INFO)` to enable it." ) if self.iteration_log and not engine.has_event_handler(self.iteration_completed, Events.ITERATION_COMPLETED): - engine.add_event_handler(Events.ITERATION_COMPLETED, self.iteration_completed) + event = Events.ITERATION_COMPLETED + if callable(self.iteration_log): # substitute event with new one using filter callable + event = event(event_filter=self.iteration_log) + engine.add_event_handler(event, self.iteration_completed) if self.epoch_log and not engine.has_event_handler(self.epoch_completed, Events.EPOCH_COMPLETED): - engine.add_event_handler(Events.EPOCH_COMPLETED, self.epoch_completed) + event = Events.EPOCH_COMPLETED + if callable(self.epoch_log): # substitute event with new one using filter callable + event = event(event_filter=self.epoch_log) + engine.add_event_handler(event, self.epoch_completed) if not engine.has_event_handler(self.exception_raised, Events.EXCEPTION_RAISED): engine.add_event_handler(Events.EXCEPTION_RAISED, self.exception_raised) diff --git a/tests/test_handler_stats.py b/tests/test_handler_stats.py index cb93f93a29..84477f9221 100644 --- a/tests/test_handler_stats.py +++ b/tests/test_handler_stats.py @@ -20,12 +20,23 @@ import torch from ignite.engine import Engine, Events +from parameterized import parameterized from monai.handlers import StatsHandler +def get_event_filter(e): + def event_filter(_, event): + if event in e: + return True + return False + + return event_filter + + class TestHandlerStats(unittest.TestCase): - def test_metrics_print(self): + @parameterized.expand([[True], [get_event_filter([1, 2])]]) + def test_metrics_print(self, epoch_log): log_stream = StringIO() log_handler = logging.StreamHandler(log_stream) log_handler.setLevel(logging.INFO) @@ -48,10 +59,11 @@ def _update_metric(engine): logger = logging.getLogger(key_to_handler) logger.setLevel(logging.INFO) logger.addHandler(log_handler) - stats_handler = StatsHandler(iteration_log=False, epoch_log=True, name=key_to_handler) + stats_handler = StatsHandler(iteration_log=False, epoch_log=epoch_log, name=key_to_handler) stats_handler.attach(engine) - engine.run(range(3), max_epochs=2) + max_epochs = 4 + engine.run(range(3), max_epochs=max_epochs) # check logging output output_str = log_stream.getvalue() @@ -61,9 +73,13 @@ def _update_metric(engine): for line in output_str.split("\n"): if has_key_word.match(line): content_count += 1 - self.assertTrue(content_count > 0) + if epoch_log is True: + self.assertTrue(content_count == max_epochs) + else: + self.assertTrue(content_count == 2) # 2 = len([1, 2]) from event_filter - def test_loss_print(self): + @parameterized.expand([[True], [get_event_filter([1, 3])]]) + def test_loss_print(self, iteration_log): log_stream = StringIO() log_handler = logging.StreamHandler(log_stream) log_handler.setLevel(logging.INFO) @@ -80,10 +96,14 @@ def _train_func(engine, batch): logger = logging.getLogger(key_to_handler) logger.setLevel(logging.INFO) logger.addHandler(log_handler) - stats_handler = StatsHandler(iteration_log=True, epoch_log=False, name=key_to_handler, tag_name=key_to_print) + stats_handler = StatsHandler( + iteration_log=iteration_log, epoch_log=False, name=key_to_handler, tag_name=key_to_print + ) stats_handler.attach(engine) - engine.run(range(3), max_epochs=2) + num_iters = 3 + max_epochs = 2 + engine.run(range(num_iters), max_epochs=max_epochs) # check logging output output_str = log_stream.getvalue() @@ -93,7 +113,10 @@ def _train_func(engine, batch): for line in output_str.split("\n"): if has_key_word.match(line): content_count += 1 - self.assertTrue(content_count > 0) + if iteration_log is True: + self.assertTrue(content_count == num_iters * max_epochs) + else: + self.assertTrue(content_count == 2) # 2 = len([1, 3]) from event_filter def test_loss_dict(self): log_stream = StringIO()