mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-01-17 16:08:30 +00:00
73d36618a6
* add memory trainer hook * fix bug * add memory trainer hook * fix import bug * fix import bug * add trainer hook * fix #370 git log bug * modify `to_tensorboard` function to support better output * remove useless output * change the name of `MemProfiler` * complete memory profiler * replace error with warning * finish trainer hook * modify interface of MemProfiler * modify `__init__.py` in profiler * remove unnecessary pass statement * add usage to doc string * add usage to trainer hook * new location to store temp data file
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
from pathlib import Path
|
|
from typing import Union
|
|
from colossalai.engine import Engine
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
from colossalai.engine.ophooks import MemTracerOpHook
|
|
from colossalai.utils.profiler import BaseProfiler
|
|
|
|
|
|
class MemProfiler(BaseProfiler):
|
|
"""Wraper of MemOpHook, used to show GPU memory usage through each iteration
|
|
|
|
To use this profiler, you need to pass an `engine` instance. And the usage is same like
|
|
CommProfiler.
|
|
|
|
mm_prof = MemProfiler(engine)
|
|
with ProfilerContext([mm_prof]) as prof:
|
|
writer = SummaryWriter("mem")
|
|
engine.train()
|
|
...
|
|
prof.to_file("./log")
|
|
prof.to_tensorboard(writer)
|
|
|
|
"""
|
|
|
|
def __init__(self, engine: Engine, warmup: int = 50, refreshrate: int = 10) -> None:
|
|
super().__init__(profiler_name="MemoryProfiler", priority=0)
|
|
self._mem_tracer = MemTracerOpHook(warmup=warmup, refreshrate=refreshrate)
|
|
self._engine = engine
|
|
|
|
def enable(self) -> None:
|
|
self._engine.add_hook(self._mem_tracer)
|
|
|
|
def disable(self) -> None:
|
|
self._engine.remove_hook(self._mem_tracer)
|
|
|
|
def to_tensorboard(self, writer: SummaryWriter) -> None:
|
|
stats = self._mem_tracer.async_mem_monitor.state_dict['mem_stats']
|
|
for info, i in enumerate(stats):
|
|
writer.add_scalar(
|
|
"memory_usage/GPU",
|
|
info,
|
|
i
|
|
)
|
|
|
|
def to_file(self, data_file: Path) -> None:
|
|
self._mem_tracer.save_results(data_file)
|
|
|
|
def show(self) -> None:
|
|
stats = self._mem_tracer.async_mem_monitor.state_dict['mem_stats']
|
|
print(stats)
|