"git@developer.sourcefind.cn:OpenDAS/pytorch3d.git" did not exist on "029a9da00b4015b6988a28e353fc389572ac6254"
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 datetime
import torch
from .base import LoggerHook from .base import LoggerHook
...@@ -41,5 +43,10 @@ class TextLoggerHook(LoggerHook): ...@@ -41,5 +43,10 @@ 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