# AUTOGENERATED FILE! PLEASE DON'T EDIT
from k1lib.callbacks import Callback, Cbs
import k1lib; from k1lib import cli
_li = 30
class IOData:
def __init__(self, ioProfiler, mS:k1lib.selector.ModuleSelector):
self.ioProfiler = ioProfiler; self.mS = mS
self.iS = None; self.oS = None
self.handle = None; self.hook()
def hook(self):
def hk(m, i, o):
self.iS = list(k1lib.squeeze(i, True).shape)
self.oS = list(k1lib.squeeze(o, True).shape)
self.handle = self.mS.nn.register_forward_hook(hk)
def unhook(self): self.handle.remove()
def __getstate__(self):
answer = dict(self.__dict__)
del answer["mS"]; del answer["ioProfiler"]; return answer
def __setstate__(self, state): self.__dict__.update(dict(state))
def __str__(self):
a = f"{self.iS}".ljust(_li); b = f"{self.oS}".ljust(_li)
return f"{a}{b}"
[docs]class IOProfiler(Callback):
"""Gets input and output shapes of each layer.
Example::
l = k1lib.Learner.sample()
l.cbs.add(Cbs.Profiler())
# views table
l.Profiler.io
# views table highlighted
l.Profiler.io.css("#lin1")
"""
def startRun(self):
if not hasattr(self, "selector"): # if no selectors found
self.selector = self.l.model.select("")
for m in self.selector.modules(): m.data = IOData(self, m)
self.selector.displayF = lambda m: (k1lib.fmt.txt.red if "_ioProf_" in m else k1lib.fmt.txt.identity)(m.data)
def startStep(self): return True
def _run(self):
"""Runs everything"""
with self.cbs.suspendEval(): self.l.run(1, 1)
for m in self.selector.modules(): m.data.unhook()
[docs] def css(self, css:str):
"""Selects a small part of the network to highlight. See also: :mod:`k1lib.selector`."""
self.selector.parse(k1lib.selector.preprocess(css, "_ioProf_"))
print(self.__repr__()); self.selector.clearProps()
def __repr__(self):
header = "input shape".ljust(_li) + "output shape".ljust(_li)
c = self.selector.__repr__(intro=False, header=header).split("\n") | cli.tab() | cli.join("\n")
return f"""IOProfiler:\n{c}
Can...
- iop.css("..."): highlights a particular part of the network
- iop.selector: to get internal k1lib.ModuleSelector object"""