Source code for k1lib.callbacks.profiler

# AUTOGENERATED FILE! PLEASE DON'T EDIT
from .callbacks import Callback, Callbacks, Cbs
import k1lib, time, torch, math, logging, numpy as np, torch.nn as nn
from functools import partial
import matplotlib.pyplot as plt
__all__ = ["Profiler"]
import k1lib.callbacks.profilers as ps
ComputationProfiler = ps.computation.ComputationProfiler
IOProfiler = ps.io.IOProfiler
MemoryProfiler = ps.memory.MemoryProfiler
TimeProfiler = ps.time.TimeProfiler
[docs]@k1lib.patch(Cbs) class Profiler(Callback): """Profiles memory, time, and computational complexity of the network. See over :mod:`k1lib.callbacks.profilers` for more details on each of these profilers""" def __init__(self): super().__init__(); self.clear(); self.dependsOn=["Recorder"]
[docs] def clear(self): """Clears every child profilers""" self._mpCache=None; self._tpCache=None self._cpCache=None; self._ioCache=None
def _memory(self): # do this to quickly debug, cause if not, Callback will just raise AttributeError on .memory if self._mpCache != None: return self._mpCache with self.cbs.context(): mp = MemoryProfiler(); self.cbs.add(mp) mp.run(); self._mpCache = mp; return mp @property def memory(self) -> MemoryProfiler: """Gets the memory profiler""" return self._memory() def _computation(self): if self._cpCache != None: return self._cpCache with self.cbs.context(): cp = ComputationProfiler(self); self.cbs.add(cp) cp.run(); self._cpCache = cp; return cp @property def computation(self) -> ComputationProfiler: """Gets the computation profiler""" return self._computation() def _time(self): if self._tpCache != None: return self._tpCache with self.cbs.context(): tp = TimeProfiler(); self.cbs.add(tp) tp.run(); self._tpCache = tp; return tp @property def time(self) -> TimeProfiler: """Gets the time profiler""" return self._time() def _io(self): if self._ioCache != None: return self._ioCache with self.cbs.context(): io = IOProfiler(); self.cbs.add(io) io.run(); self._ioCache = io; return io @property def io(self) -> IOProfiler: """Gets the IO profiler""" return self._io() def __repr__(self): return f"""{self._reprHead}, can... - p.memory: to profile module memory requirements - p.time: to profile module execution times - p.computation: to estimate module computation - p.io: to get input and output shapes of {self._reprCan}"""