|
|
|
import functools |
|
|
|
|
|
class OutputHook: |
|
|
|
def __init__(self, module, outputs=None, as_tensor=False): |
|
self.outputs = outputs |
|
self.as_tensor = as_tensor |
|
self.layer_outputs = {} |
|
self.register(module) |
|
|
|
def register(self, module): |
|
|
|
def hook_wrapper(name): |
|
|
|
def hook(model, input, output): |
|
if self.as_tensor: |
|
self.layer_outputs[name] = output |
|
else: |
|
if isinstance(output, list): |
|
self.layer_outputs[name] = [ |
|
out.detach().cpu().numpy() for out in output |
|
] |
|
else: |
|
self.layer_outputs[name] = output.detach().cpu().numpy( |
|
) |
|
|
|
return hook |
|
|
|
self.handles = [] |
|
if isinstance(self.outputs, (list, tuple)): |
|
for name in self.outputs: |
|
try: |
|
layer = rgetattr(module, name) |
|
h = layer.register_forward_hook(hook_wrapper(name)) |
|
except ModuleNotFoundError as module_not_found: |
|
raise ModuleNotFoundError( |
|
f'Module {name} not found') from module_not_found |
|
self.handles.append(h) |
|
|
|
def remove(self): |
|
for h in self.handles: |
|
h.remove() |
|
|
|
def __enter__(self): |
|
return self |
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb): |
|
self.remove() |
|
|
|
|
|
|
|
|
|
def rgetattr(obj, attr, *args): |
|
|
|
def _getattr(obj, attr): |
|
return getattr(obj, attr, *args) |
|
|
|
return functools.reduce(_getattr, [obj] + attr.split('.')) |
|
|
