Commit 0f88eddb authored by wangg12's avatar wangg12
Browse files

allow passing custom logger to runner

parent ddb52228
...@@ -27,6 +27,7 @@ class Runner(object): ...@@ -27,6 +27,7 @@ class Runner(object):
work_dir (str, optional): The working directory to save checkpoints work_dir (str, optional): The working directory to save checkpoints
and logs. and logs.
log_level (int): Logging level. log_level (int): Logging level.
logger: Custom logger. If `None`, use default logger.
""" """
def __init__(self, def __init__(self,
...@@ -34,7 +35,8 @@ class Runner(object): ...@@ -34,7 +35,8 @@ class Runner(object):
batch_processor, batch_processor,
optimizer=None, optimizer=None,
work_dir=None, work_dir=None,
log_level=logging.INFO): log_level=logging.INFO,
logger=None):
assert callable(batch_processor) assert callable(batch_processor)
self.model = model self.model = model
if optimizer is not None: if optimizer is not None:
...@@ -59,7 +61,10 @@ class Runner(object): ...@@ -59,7 +61,10 @@ class Runner(object):
self._model_name = self.model.__class__.__name__ self._model_name = self.model.__class__.__name__
self._rank, self._world_size = get_dist_info() self._rank, self._world_size = get_dist_info()
self.logger = self.init_logger(work_dir, log_level) if logger is None:
self.logger = self.init_logger(work_dir, log_level)
else:
self.logger = logger
self.log_buffer = LogBuffer() self.log_buffer = LogBuffer()
self.mode = None self.mode = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment