Commit f3d5d979 authored by Cao Yuhang's avatar Cao Yuhang
Browse files

fix memory order in log text

parent 1263c91e
...@@ -36,6 +36,12 @@ class TextLoggerHook(LoggerHook): ...@@ -36,6 +36,12 @@ class TextLoggerHook(LoggerHook):
log_str += ( log_str += (
'time: {log[time]:.3f}, data_time: {log[data_time]:.3f}, '. 'time: {log[time]:.3f}, data_time: {log[data_time]:.3f}, '.
format(log=runner.log_buffer.output)) format(log=runner.log_buffer.output))
# statistic memory
if runner.mode == 'train' and torch.cuda.is_available():
mem = torch.cuda.max_memory_allocated()
mem_mb = int(mem / (1024 * 1024))
mem_str = 'memory: {}, '.format(mem_mb)
log_str += mem_str
log_items = [] log_items = []
for name, val in runner.log_buffer.output.items(): for name, val in runner.log_buffer.output.items():
if name in ['time', 'data_time']: if name in ['time', 'data_time']:
...@@ -43,10 +49,5 @@ class TextLoggerHook(LoggerHook): ...@@ -43,10 +49,5 @@ class TextLoggerHook(LoggerHook):
if isinstance(val, float): if isinstance(val, float):
val = '{:.4f}'.format(val) val = '{:.4f}'.format(val)
log_items.append('{}: {}'.format(name, val)) log_items.append('{}: {}'.format(name, val))
# statistic memory
if runner.mode == 'train' and torch.cuda.is_available():
mem = torch.cuda.max_memory_allocated()
mem_mb = int(mem / (1024 * 1024))
log_items.append('{}: {}'.format('memory', mem_mb))
log_str += ', '.join(log_items) log_str += ', '.join(log_items)
runner.logger.info(log_str) runner.logger.info(log_str)
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