neural-feature-importance 0.5.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

neural_feature_importance/__init__.py

@@ -7,6 +7,11 @@ from .callbacks import (
     VarianceImportanceKeras,
     VarianceImportanceTorch,
 )
+from .conv_callbacks import ConvVarianceImportanceKeras, ConvVarianceImportanceTorch
+from .embedding_callbacks import (
+    EmbeddingVarianceImportanceKeras,
+    EmbeddingVarianceImportanceTorch,
+)
 from .utils import MetricThreshold
 
 try:
@@ -19,4 +24,8 @@ __all__ = [
     "VarianceImportanceKeras",
     "VarianceImportanceTorch",
     "MetricThreshold",
+    "ConvVarianceImportanceKeras",
+    "ConvVarianceImportanceTorch",
+    "EmbeddingVarianceImportanceKeras",
+    "EmbeddingVarianceImportanceTorch",
 ]

neural_feature_importance/callbacks.py

@@ -1,4 +1,10 @@
-"""Variance-based feature importance utilities."""
+"""Utilities for computing variance-based feature importances.
+
+These classes track the weights of the first trainable layer during training
+and estimate feature importances by accumulating the variance of each weight
+value. After training, the variances are combined with the last observed
+weights to produce a normalized importance score for every input feature.
+"""
 
 from __future__ import annotations
 
@@ -13,7 +19,13 @@ logger = logging.getLogger(__name__)
 
 
 class VarianceImportanceBase:
-    """Compute feature importance using Welford's algorithm."""
+    """Compute feature importances using running variance statistics.
+
+    The class implements Welford's algorithm to accumulate the variance of
+    weight values over training iterations. Feature importances are derived by
+    combining the final variance estimates with the absolute value of the last
+    observed weights.
+    """
 
     def __init__(self) -> None:
         self._n = 0
@@ -23,13 +35,21 @@ class VarianceImportanceBase:
         self.var_scores: np.ndarray | None = None
 
     def start(self, weights: np.ndarray) -> None:
-        """Initialize statistics for the given weight matrix."""
+        """Initialize running statistics.
+
+        Parameters
+        ----------
+        weights:
+            Initial weight matrix of shape ``(features, outputs)``. The values
+            are converted to ``float64`` for numerical stability and the running
+            mean and variance buffers are reset.
+        """
         self._mean = weights.astype(np.float64)
         self._m2 = np.zeros_like(self._mean)
         self._n = 0
 
     def update(self, weights: np.ndarray) -> None:
-        """Update running statistics with new weights."""
+        """Update running mean and variance using new weights."""
         if self._mean is None or self._m2 is None:
             return
         self._n += 1
@@ -40,7 +60,7 @@ class VarianceImportanceBase:
         self._last_weights = weights
 
     def finalize(self) -> None:
-        """Finalize statistics and compute normalized scores."""
+        """Compute normalized importance scores from accumulated statistics."""
         if self._last_weights is None or self._m2 is None:
             logger.warning(
                 "%s was not fully initialized; no scores computed", self.__class__.__name__
@@ -67,6 +87,7 @@ class VarianceImportanceBase:
         return self.var_scores
 
 
+
 class VarianceImportanceKeras(Callback, VarianceImportanceBase):
     """Keras callback implementing variance-based feature importance."""
 

neural_feature_importance/conv_callbacks.py

@@ -0,0 +1,104 @@
+"""Callbacks that extend variance tracking to convolutional layers."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import numpy as np
+
+from .callbacks import VarianceImportanceKeras, VarianceImportanceTorch
+
+logger = logging.getLogger(__name__)
+
+
+def _flatten_weights(weights: np.ndarray, outputs_last: bool) -> np.ndarray:
+    """Return a two-dimensional view of convolutional kernels.
+
+    Parameters
+    ----------
+    weights:
+        Weight tensor from a convolutional layer. Expected shape is
+        ``(H, W, in_channels, out_channels)`` when ``outputs_last`` is ``True``
+        and ``(out_channels, in_channels, H, W)`` otherwise.
+    outputs_last:
+        Whether the output dimension is the last axis of ``weights``.
+
+    Returns
+    -------
+    np.ndarray
+        Array of shape ``(features, outputs)`` suitable for variance tracking.
+    """
+    if weights.ndim > 2:
+        if outputs_last:
+            return weights.reshape(-1, weights.shape[-1])
+        return weights.reshape(weights.shape[0], -1).T
+    return weights
+
+
+class ConvVarianceImportanceKeras(VarianceImportanceKeras):
+    """Keras callback that tracks convolutional kernels.
+
+    The first trainable layer is inspected and, if its weights have more than
+    two dimensions, they are flattened so that each spatial location and input
+    channel is treated as a separate feature. Variances are accumulated during
+    training and converted to per-filter importance scores.
+    """
+    def on_train_begin(self, logs: Optional[dict] = None) -> None:
+        self._layer = None
+        for layer in self.model.layers:
+            has_vars = bool(layer.trainable_weights)
+            has_data = bool(layer.get_weights())
+            if has_vars and has_data:
+                self._layer = layer
+                break
+        if self._layer is None:
+            raise ValueError("Model does not contain trainable weights.")
+        weights = self._layer.get_weights()[0]
+        weights = _flatten_weights(weights, outputs_last=True)
+        logger.info(
+            "Tracking variance for layer '%s' with %d features",
+            self._layer.name,
+            weights.shape[0],
+        )
+        self.start(weights)
+
+    def on_epoch_end(self, epoch: int, logs: Optional[dict] = None) -> None:
+        if self._layer is None:
+            return
+        weights = self._layer.get_weights()[0]
+        weights = _flatten_weights(weights, outputs_last=True)
+        self.update(weights)
+
+
+class ConvVarianceImportanceTorch(VarianceImportanceTorch):
+    """PyTorch helper with convolutional support.
+
+    Works analogously to :class:`ConvVarianceImportanceKeras` but for models
+    built with :mod:`torch.nn`. The first trainable parameter with two or more
+    dimensions is flattened so each spatial position becomes a tracked feature.
+    """
+    def on_train_begin(self) -> None:
+        from torch import nn
+
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and param.dim() >= 2:
+                self._param = param
+                weights = param.detach().cpu().numpy()
+                weights = _flatten_weights(weights, outputs_last=False)
+                logger.info(
+                    "Tracking variance for parameter '%s' with %d features",
+                    name,
+                    weights.shape[0],
+                )
+                self.start(weights)
+                break
+        if self._param is None:
+            raise ValueError("Model does not contain trainable parameters")
+
+    def on_epoch_end(self) -> None:
+        if self._param is None:
+            return
+        weights = self._param.detach().cpu().numpy()
+        weights = _flatten_weights(weights, outputs_last=False)
+        self.update(weights)

neural_feature_importance/embedding_callbacks.py

@@ -0,0 +1,85 @@
+"""Callbacks that compute variance-based importance for embedding layers.
+
+These callbacks extend :class:`~neural_feature_importance.callbacks.VarianceImportanceBase`
+to operate on 2-D embedding matrices. The variance of each embedding vector is
+accumulated over training and the resulting per-token scores are normalized
+between 0 and 1.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+from .callbacks import VarianceImportanceKeras, VarianceImportanceTorch
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingVarianceImportanceKeras(VarianceImportanceKeras):
+    """Variance-based importance callback for Keras embedding layers.
+
+    During training this callback monitors the weights of the first trainable
+    layer (expected to be an :class:`~tensorflow.keras.layers.Embedding`) and
+    accumulates the running variance of each embedding vector. After training the
+    variances are summed across the embedding dimension to yield a single score
+    per token.
+    """
+
+    def finalize(self) -> None:  # type: ignore[override]
+        if self._last_weights is None or self._m2 is None:
+            logger.warning(
+                "%s was not fully initialized; no scores computed",
+                self.__class__.__name__,
+            )
+            return
+
+        if self._n < 2:
+            variance = np.full_like(self._m2, np.nan)
+        else:
+            variance = self._m2 / (self._n - 1)
+
+        scores = np.sum(variance, axis=1)
+        min_val = float(np.nanmin(scores))
+        max_val = float(np.nanmax(scores))
+        denom = max_val - min_val if max_val != min_val else 1.0
+        self.var_scores = (scores - min_val) / denom
+
+        top = np.argsort(self.var_scores)[-10:][::-1]
+        logger.info("Most important tokens: %s", top)
+
+
+class EmbeddingVarianceImportanceTorch(VarianceImportanceTorch):
+    """Variance-based importance for PyTorch embedding layers.
+
+    Parameters
+    ----------
+    model:
+        Neural network containing an :class:`torch.nn.Embedding` layer whose
+        weights will be monitored.
+    """
+
+    def finalize(self) -> None:  # type: ignore[override]
+        if self._last_weights is None or self._m2 is None:
+            logger.warning(
+                "%s was not fully initialized; no scores computed",
+                self.__class__.__name__,
+            )
+            return
+
+        if self._n < 2:
+            variance = np.full_like(self._m2, np.nan)
+        else:
+            variance = self._m2 / (self._n - 1)
+
+        scores = np.sum(variance, axis=1)
+        min_val = float(np.nanmin(scores))
+        max_val = float(np.nanmax(scores))
+        denom = max_val - min_val if max_val != min_val else 1.0
+        self.var_scores = (scores - min_val) / denom
+
+        top = np.argsort(self.var_scores)[-10:][::-1]
+        logger.info("Most important tokens: %s", top)
+
+

{neural_feature_importance-0.5.2.dist-info → neural_feature_importance-0.9.1.dist-info}/METADATA

@@ -1,20 +1,23 @@
 Metadata-Version: 2.4
 Name: neural-feature-importance
-Version: 0.5.2
+Version: 0.9.1
 Summary: Variance-based feature importance for Neural Networks using callbacks for Keras and PyTorch
 Author: CR de Sá
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: numpy
 Provides-Extra: tensorflow
 Requires-Dist: tensorflow; extra == "tensorflow"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
+Dynamic: license-file
 
 # neural-feature-importance
 
 [![PyPI version](https://img.shields.io/pypi/v/neural-feature-importance.svg)](https://pypi.org/project/neural-feature-importance/)
-[![Python versions](https://img.shields.io/pypi/pyversions/neural-feature-importance.svg)](https://pypi.org/project/neural-feature-importance/)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)
 
 Variance-based feature importance for deep learning models.
 
@@ -74,19 +77,44 @@ print(tracker.feature_importances_)
 
 ## Example scripts
 
-Run `compare_feature_importance.py` to train a small network on the Iris dataset
+Run `scripts/compare_feature_importance.py` to train a small network on the Iris dataset
 and compare the scores with a random forest baseline:
 
 ```bash
 python compare_feature_importance.py
 ```
 
-Run `full_experiment.py` to reproduce the experiments from the paper:
+Run `scripts/full_experiment.py` to reproduce the experiments from the paper:
 
 ```bash
 python full_experiment.py
 ```
 
+### Convolutional models
+
+To compute importances for convolutional networks, use
+`ConvVarianceImportanceKeras` from `neural_feature_importance.conv_callbacks`.
+`scripts/conv_visualization_example.py` trains small Conv2D models on the MNIST
+and scikit‑learn digits datasets and displays per-filter heatmaps. An equivalent
+notebook is available in ``notebooks/conv_visualization_example.ipynb``:
+
+```bash
+python scripts/conv_visualization_example.py
+```
+
+### Embedding layers
+
+To compute token importances from embedding weights, use
+`EmbeddingVarianceImportanceKeras` or `EmbeddingVarianceImportanceTorch` from
+`neural_feature_importance.embedding_callbacks`.
+Run `scripts/token_importance_topk_example.py` to train a small text classifier
+on IMDB and display the most important tokens. A matching notebook lives in
+``notebooks/token_importance_topk_example.ipynb``:
+
+```bash
+python scripts/token_importance_topk_example.py
+```
+
 ## Development
 
 After making changes, run the following checks:
@@ -124,3 +152,7 @@ If you use this package in your research, please cite:
 ```
 
 We appreciate citations as they help the community discover this work.
+
+## License
+
+This project is licensed under the [MIT License](LICENSE).

neural_feature_importance-0.9.1.dist-info/RECORD

@@ -0,0 +1,11 @@
+neural_feature_importance/__init__.py,sha256=EICAFjL6mHquX5wd1YWeV_6hI_jTgqiGNvKjYL7URSU,914
+neural_feature_importance/callbacks.py,sha256=dPrxkjh6inf8hI8wGhNL-elBjWuCEWQHWICNmTYqbyE,5677
+neural_feature_importance/conv_callbacks.py,sha256=L1u7EVAERtRBpNlWcH3u5A49jBcALe8Yv0giMxTdKPM,3789
+neural_feature_importance/embedding_callbacks.py,sha256=TYc4Xu2MzK3Ff0JBFakiEoZtzCbRsVjoswPocheFcr0,2869
+neural_feature_importance/utils/__init__.py,sha256=dMjBUCx8DCoJKAEAnjj_daXfEu9Q5va1k8XupmWdZiE,114
+neural_feature_importance/utils/monitors.py,sha256=LTz7oE0-WgZ50DHyHDnTwfzWSSWMnjWd0xlwt7BWKuU,1763
+neural_feature_importance-0.9.1.dist-info/licenses/LICENSE,sha256=6v0bh8lk889d7vmcFAqzUbqly-ogYCYdcCcTW4yZ2tg,1066
+neural_feature_importance-0.9.1.dist-info/METADATA,sha256=l2Q2CbawJVfsux8Y80Vx2bJ_FC72SH-9tLrXy9dpw50,5187
+neural_feature_importance-0.9.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+neural_feature_importance-0.9.1.dist-info/top_level.txt,sha256=yP0Q-BG7hDLLu1H1_x5bGEKwkCso5NxxvScnlmICb-o,26
+neural_feature_importance-0.9.1.dist-info/RECORD,,

neural_feature_importance-0.9.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 CR de Sá
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

neural_feature_importance-0.5.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-neural_feature_importance/__init__.py,sha256=z3Rve0a7QTAhEpCesDhSdbkOwfSsNZgiwDVep1Is_c0,566
-neural_feature_importance/callbacks.py,sha256=HMHsmVaqZOzy5NSbxN-8CWvq82vzgZgZD53zqp2nAz0,4811
-neural_feature_importance/utils/__init__.py,sha256=dMjBUCx8DCoJKAEAnjj_daXfEu9Q5va1k8XupmWdZiE,114
-neural_feature_importance/utils/monitors.py,sha256=LTz7oE0-WgZ50DHyHDnTwfzWSSWMnjWd0xlwt7BWKuU,1763
-neural_feature_importance-0.5.2.dist-info/METADATA,sha256=W5HuRD-lpppfwv8DhQ26eevxiZUkL-vaevfQDLdbKak,4091
-neural_feature_importance-0.5.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-neural_feature_importance-0.5.2.dist-info/top_level.txt,sha256=yP0Q-BG7hDLLu1H1_x5bGEKwkCso5NxxvScnlmICb-o,26
-neural_feature_importance-0.5.2.dist-info/RECORD,,