Source code for k1lib.callbacks.profiler

# AUTOGENERATED FILE! PLEASE DON'T EDIT
from k1lib.callbacks import Callback, Callbacks
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(Callback.cls) 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.append(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.append(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.append(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.append(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}"""
@k1lib.patch(Callbacks, docs=Profiler) def withProfiler(self): return self.append(Profiler())