| | |
| | |
| | |
| | import logging |
| | import torch |
| | from .functions import REG_FUNCTION_MAP |
| |
|
| |
|
| | |
| | |
| | |
| | class HookMonitor: |
| | """ |
| | Monitors forward activations and backward gradients of a PyTorch model by |
| | registering hooks on all its submodules. The monitor computes per-layer |
| | statistics defined in `REG_FUNCTION_MAP`, accumulating them during forward |
| | and backward passes, and provides normalized results at the end. |
| | |
| | This class is designed to be lightweight, safe (uses no_grad for activation |
| | hooks), and usable as a context manager to automate attachment and cleanup |
| | of hooks. |
| | |
| | ---------------------------------------- |
| | Core Behavior |
| | ---------------------------------------- |
| | - During the forward pass: |
| | • A forward hook receives (module, input, output). |
| | • The activation tensor is detached and cast to float. |
| | • For each registered metric in REG_FUNCTION_MAP, if its watcher flag |
| | is enabled, the metric is computed and accumulated. |
| | • A gradient hook is registered on the output tensor so that gradient |
| | statistics can also be collected during backpropagation. |
| | |
| | - During backpropagation: |
| | • The gradient hook receives the gradient tensor for the activation. |
| | • Any metric marked as `grad_<metric>` in the watcher dictionary will be |
| | applied to the gradient tensor and accumulated. |
| | |
| | - Statistics: |
| | • For each metric, the class tracks both the accumulated value and a |
| | "/valid/" counter. |
| | • `get_stats()` returns normalized statistics (sum / valid_count) for |
| | each metric per layer. |
| | |
| | ---------------------------------------- |
| | Parameters |
| | ---------------------------------------- |
| | model : torch.nn.Module |
| | The model whose modules will be monitored. All submodules returned by |
| | `model.named_modules()` will receive a forward hook. |
| | |
| | watcher : dict |
| | A dictionary mapping metric names to boolean flags. Keys must match the |
| | names used in `REG_FUNCTION_MAP`. Example: |
| | { |
| | "mean": True, |
| | "std": True, |
| | "grad_mean": True |
| | } |
| | |
| | Metrics not enabled here will not be computed. |
| | |
| | logger : logging.Logger |
| | A Logger used to report errors, debugging information, and warnings. |
| | |
| | ---------------------------------------- |
| | Attributes |
| | ---------------------------------------- |
| | stats : dict |
| | Nested dictionary storing accumulated statistics per layer. Normalized |
| | results are returned by `get_stats()`. |
| | |
| | handles : list |
| | A List of hook handles returned by `register_forward_hook`. These are |
| | stored to later remove all hooks safely. |
| | |
| | ---------------------------------------- |
| | Usage Example |
| | ---------------------------------------- |
| | >>> model: torch.nn.Module |
| | >>> watcher: dict[str, bool] |
| | >>> logger: logging.Logger |
| | >>> x: torch.Tensor |
| | >>> loss: torch.nn.Module # Loss |
| | |
| | >>> monitor = HookMonitor(model, watcher, logger) |
| | >>> monitor.attach() |
| | >>> output = model(x) |
| | >>> loss.backward() |
| | >>> stats = monitor.get_stats() |
| | >>> monitor.remove() |
| | |
| | Or using a context manager: |
| | |
| | >>> with HookMonitor(model, watcher, logger) as monitor: |
| | ... output = model(x) |
| | ... loss.backward() |
| | >>> stats = monitor.get_stats() |
| | |
| | ---------------------------------------- |
| | Notes |
| | ---------------------------------------- |
| | - The gradient hook is attached to the activation tensor (module output), |
| | not to model parameters. |
| | - No gradients are tracked during forward hooks thanks to @torch.no_grad(). |
| | - The monitor does not interfere with the training process: it only reads |
| | activations and gradients. |
| | - Missing '/valid/' counters trigger an error log and skip normalization for |
| | that metric. |
| | |
| | """ |
| | def __init__(self, model: torch.nn.Module, watcher: dict, logger: logging.Logger): |
| | """ |
| | Initialize a HookMonitor instance to track activation and gradient |
| | statistics across all modules of a PyTorch model. |
| | |
| | This constructor does not attach any hooks yet; it simply stores the |
| | monitoring configuration. Hooks are registered only when `attach()` or |
| | the context manager (`with HookMonitor(...)`) is used. |
| | |
| | Parameters |
| | ---------- |
| | model : torch.nn.Module |
| | The model whose internal modules will be monitored. Every submodule |
| | returned by `model.named_modules()` will receive a forward hook. |
| | |
| | watcher : dict |
| | Dictionary of boolean flags controlling which statistics should be |
| | computed. Keys must match the names in `REG_FUNCTION_MAP`. |
| | Example: |
| | { |
| | "mean": True, |
| | "std": False, |
| | "grad_mean": True |
| | } |
| | |
| | Any metric not enabled here will not be computed during execution. |
| | |
| | logger : logging.Logger |
| | Logging instance used for reporting errors, debug messages and |
| | warnings during monitoring operations. |
| | |
| | Attributes Initialized |
| | ---------------------- |
| | model : torch.nn.Module |
| | Stored reference to the monitored model. |
| | |
| | watcher : dict |
| | The watcher configuration controlling metric activation. |
| | |
| | stats : dict |
| | Internal dictionary used to accumulate statistics across all layers. |
| | |
| | handles : list |
| | A List of hook handles created when calling `.attach()`. Each handle |
| | is later used to safely remove hooks with `.remove()`. |
| | |
| | Notes |
| | ----- |
| | - No hooks are installed at construction time. |
| | - The monitor becomes active only after calling `.attach()` or entering |
| | a `with` block. |
| | """ |
| | self.logger: logging.Logger = logger |
| | self.model: torch.nn.Module = model |
| | self.watcher: dict = watcher |
| | self.stats: dict = dict() |
| | self.handles: list = list() |
| |
|
| | def _build_hook(self, name): |
| |
|
| | @torch.no_grad() |
| | def hook(*args): |
| | _, _, act = args |
| |
|
| | if torch.is_tensor(act): |
| | act_detached = act.detach().float() |
| | s = self.stats.setdefault(name, {}) |
| |
|
| | |
| | for function_name, compute_function in REG_FUNCTION_MAP.items(): |
| | if self.watcher.get(function_name, False) and not function_name.startswith('grad_'): |
| | value = compute_function(act_detached, ...) |
| | if value is not None: |
| | s[function_name] = s.get(function_name, 0.0) + value |
| | s[function_name + '/valid/'] = s.get(function_name + '/valid/', 0.0) + 1 |
| |
|
| | |
| | def grad_hook(grad): |
| | gd = grad.detach().float() |
| | |
| | for gd_function_name, gd_compute_function in REG_FUNCTION_MAP.items(): |
| | if self.watcher.get('grad_' + gd_function_name, False) and not gd_function_name.startswith('grad_'): |
| | gd_function_name = 'grad_' + gd_function_name |
| | gd_value = gd_compute_function(gd, ...) |
| | if gd_value is not None: |
| | s[gd_function_name] = s.get(gd_function_name, 0.0) + gd_value |
| | s[gd_function_name + '/valid/'] = s.get(gd_function_name + '/valid/', 0.0) + 1 |
| |
|
| | if act.requires_grad: |
| | act.register_hook(grad_hook) |
| |
|
| | return hook |
| |
|
| | def get_stats(self) -> dict: |
| | """ |
| | Get the statistics of the hooks. |
| | :return: A dictionary with the statistics. |
| | """ |
| | stats = dict() |
| | for layer_name, layer_stats in self.stats.items(): |
| | sub_stats = dict() |
| | for key, item in layer_stats.items(): |
| | if '/valid/' not in key: |
| | if key + '/valid/' in layer_stats: |
| | sub_stats[key] = item / layer_stats[key + '/valid/'] |
| | else: |
| | self.logger.error(f"Key {key} has no valid count, skipping normalization.") |
| | sub_stats[key] = item |
| | stats[layer_name] = sub_stats |
| | return stats |
| |
|
| | def attach(self): |
| | """ |
| | Registers all the hooks in the model. |
| | :return: The object. |
| | """ |
| | for name, module in self.model.named_modules(): |
| | h = module.register_forward_hook(self._build_hook(name)) |
| | self.handles.append(h) |
| | return self |
| |
|
| | def clear(self): |
| | """ |
| | Clear stats' dictionary. |
| | :return: Nothing |
| | """ |
| | self.stats.clear() |
| |
|
| | def remove(self): |
| | """ |
| | Remove all the hooks from the model. |
| | :return: Nothing. |
| | """ |
| | for h in self.handles: |
| | h.remove() |
| | self.handles.clear() |
| |
|
| | def __enter__(self): |
| | self.logger.debug("[Hooks] Attaching HookMonitor...") |
| | return self.attach() |
| |
|
| | def __exit__(self, exc_type, exc_val, exc_tb): |
| | self.logger.debug("[Hooks] Removing HookMonitor...") |
| | self.remove() |
| |
|
| | |
| | |
| | |
| |
|
