Commit 1263c91e authored by Cao Yuhang's avatar Cao Yuhang Committed by Kai Chen
Browse files

add memory statistic (#41)

* add memory statistic

* fix import order and cuda check
parent a0ce998e
import datetime
import torch
from .base import LoggerHook
......@@ -41,5 +43,10 @@ class TextLoggerHook(LoggerHook):
if isinstance(val, float):
val = '{:.4f}'.format(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)
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