Add utility to compute sparsity in each network type.

Author: mitkonikov

README.md

@@ -55,6 +55,9 @@ There are 3 different types of Forward-Forward-based Neural Networks implemented
  - [**FF+C**](#ffc) - [example usage](./examples/ff_c_mnist.py)
  - [**FFRNN**](#ffrnn) - [example usage](./examples/ff_rnn_mnist.py)
 
+ > [!NOTE]
+ > There's also a [backpropagation example](./examples/bp_mnist.py) implemented as a base case.
+
 ### FFNet
 
 The basic example of a Neural Network based on the Forward-Forward Algorithm.
@@ -92,6 +95,22 @@ from the previous layer, but also backward weights from the next layer.
 These networks have to be trained with multiple frames per batch, thus
 requiring even more time for both, training and inference.
 
+## Utilities
+
+We aim to implement many small utilities for training, validating, testing, developing, debugging and analysing FF networks.
+For each type of network, we have build small network-specifc train-test suites that allow you to quickly train and test
+a FF network with some specific [DataProcessor](./src/fflib/utils/data/dataprocessor.py).
+For each new dataset, you just need to define a data processor.
+The suite will take in the network and the data processor and train the network in the way it is supposed to.
+There are currently implemented 4 different testing suites:
+ - [FF Suite](./src/fflib/utils/ff_suite.py)
+ - [FF+C Suite](./src/fflib/utils/ffc_suite.py)
+ - [FFRNN Suite](./src/fflib/utils/ffrnn_suite.py)
+ - [BP Suite](./src/fflib/utils/bp_suite.py)
+
+Furthermore, each network implements some analysis functions, currently only a `sparsity` function is implemented.
+Sparsity can be computed in two ways `HOYER` [3] and `ENTROPY_BASED`.
+
 ## Contributions
 
 We really appreciate contributions from the community!
@@ -151,3 +170,4 @@ Here are a few guidelines to following while contributing on the library:
 
  - [**[1]**](https://arxiv.org/abs/2212.13345) - Hinton, G. (2022). The Forward-Forward Algorithm: Some preliminary investigations.
  - [**[2]**](https://pytorch.org/) - PyTorch.
+ - [**[3]**](https://dl.acm.org/doi/10.5555/1005332.1044709) - Hoyer, P. (2004). Non-negative Matrix Factorization with Sparseness Constraints.

src/fflib/enums.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class SparsityType(Enum):
+    HOYER = 1
+    ENTROPY_BASED = 2

src/fflib/interfaces/iff_recurrent_layer.py

@@ -1,8 +1,9 @@
 import torch
 
 from torch.nn import Module
+from fflib.enums import SparsityType
 from abc import ABC, abstractmethod
-from typing import Callable, Tuple, List
+from typing import Callable, Tuple, List, Dict
 
 
 class IFFRecurrentLayer(ABC, Module):
@@ -52,3 +53,7 @@ def run_train(
     @abstractmethod
     def strip_down(self) -> None:
         pass
+
+    @abstractmethod
+    def sparsity(self, type: SparsityType) -> Dict[str, float]:
+        pass

src/fflib/nn/ff_linear.py

@@ -2,6 +2,9 @@
 from torch.nn import Linear, Module, ReLU
 from torch.optim import Adam, Optimizer
 
+from fflib.enums import SparsityType
+from fflib.utils.maths import ComputeSparsity
+
 from typing import Callable, Tuple, Any, cast
 
 
@@ -131,3 +134,7 @@ def run_train(self, x_pos: torch.Tensor, x_neg: torch.Tensor) -> None:
 
     def strip_down(self) -> None:
         self.opt = None
+
+    def sparsity(self, type: SparsityType) -> torch.Tensor:
+        """Computes the sparsity of the weight's matrix."""
+        return ComputeSparsity(torch.flatten(self.weight), type)

src/fflib/nn/ff_net.py

@@ -3,6 +3,7 @@
 from torch.nn import Module
 from fflib.interfaces.iff import IFF
 from fflib.nn.ff_linear import FFLinear
+from fflib.enums import SparsityType
 from typing import List, Any, Dict, Callable
 
 
@@ -68,3 +69,8 @@ def strip_down(self) -> None:
         for layer in self.layers:
             layer.strip_down()
         delattr(self, "hooks")
+
+    def sparsity(self, type: SparsityType) -> Dict[str, float]:
+        return {
+            f"layer_{i}": float(layer.sparsity(type).item()) for i, layer in enumerate(self.layers)
+        }

src/fflib/nn/ff_recurrent_layer.py

@@ -4,7 +4,9 @@
 from torch.nn import Module, ReLU
 from torch.optim import Adam, Optimizer
 from fflib.interfaces.iff_recurrent_layer import IFFRecurrentLayer
-from typing import Callable, List, Tuple, cast, Any
+from fflib.enums import SparsityType
+from fflib.utils.maths import ComputeSparsity
+from typing import Callable, List, Tuple, Dict, cast, Any
 
 
 class FFRecurrentLayer(IFFRecurrentLayer):
@@ -144,6 +146,17 @@ def run_train(
     def strip_down(self) -> None:
         self.opt = None
 
+    def sparsity(self, type: SparsityType) -> Dict[str, float]:
+        return {
+            "fw": float(ComputeSparsity(torch.flatten(self.fw), type).item()),
+            "bw": float(ComputeSparsity(torch.flatten(self.bw), type).item()),
+            "fw+bw": float(
+                ComputeSparsity(
+                    torch.cat((torch.flatten(self.fw), torch.flatten(self.fb))), type
+                ).item()
+            ),
+        }
+
 
 class FFRecurrentLayerDummy(IFFRecurrentLayer):
     def __init__(self, dimensions: int):
@@ -175,3 +188,6 @@ def run_train(
 
     def strip_down(self) -> None:
         pass
+
+    def sparsity(self, type: SparsityType) -> Dict[str, float]:
+        return {}

src/fflib/nn/ff_rnn.py

@@ -5,7 +5,8 @@
 from fflib.nn.ff_recurrent_layer import FFRecurrentLayer, FFRecurrentLayerDummy
 from fflib.interfaces.iff import IFF
 from fflib.interfaces.iff_recurrent_layer import IFFRecurrentLayer
-from typing import List, Tuple, cast, Any
+from fflib.enums import SparsityType
+from typing import List, Tuple, Dict, cast, Any
 from typing_extensions import Self
 
 
@@ -231,3 +232,7 @@ def strip_down(self) -> None:
         for layer in self.layers:
             layer.strip_down()
         delattr(self, "hooks")
+
+    def sparsity(self, type: SparsityType) -> Dict[str, Dict[str, float]]:
+        """Returns a dictionary of dictionaries describing the sparsity levels at each layer."""
+        return {f"layer_{i}": layer.sparsity(type) for i, layer in enumerate(self.layers)}

src/fflib/nn/ffc.py

@@ -4,8 +4,10 @@
 from torch.optim import Adam, Optimizer
 from fflib.interfaces.iff import IFF
 from fflib.nn.ff_linear import FFLinear
+from fflib.enums import SparsityType
+from fflib.utils.maths import ComputeSparsity
 
-from typing import List, Callable, Any
+from typing import List, Dict, Callable, Any
 
 
 class FFC(IFF, Module):
@@ -90,3 +92,12 @@ def strip_down(self) -> None:
             layer.strip_down()
         self.optimizer = None
         delattr(self, "hooks")
+
+    def sparsity(self, type: SparsityType) -> Dict[str, float]:
+        result = {
+            f"layer_{i}": float(layer.sparsity(type).item()) for i, layer in enumerate(self.layers)
+        }
+        result["classifier"] = float(
+            ComputeSparsity(torch.flatten(self.classifier.weight), type).item()
+        )
+        return result

src/fflib/utils/maths.py

@@ -0,0 +1,19 @@
+import torch
+
+from fflib.enums import SparsityType
+
+from math import sqrt, log2
+
+
+def ComputeSparsity(x: torch.Tensor, type: SparsityType) -> torch.Tensor:
+    n = x.shape[0]
+    if type == SparsityType.HOYER:
+        r = torch.sum(torch.abs(x)) / torch.sqrt(torch.sum(torch.square(x)))
+        sqn = sqrt(n)
+        return torch.Tensor((sqn - r) / (sqn - 1))
+    elif type == SparsityType.ENTROPY_BASED:
+        t = torch.sum(torch.abs(x))
+        p = torch.abs(x).div(t) + 1e-8
+        v = p * torch.log2(p)
+        h = -torch.sum(v)
+        return torch.Tensor(1 - (h / log2(n)))

tests/nn/test_ff_linear.py

@@ -1,7 +1,7 @@
-import pytest
 import torch
 
 from fflib.nn.ff_linear import FFLinear
+from fflib.enums import SparsityType
 
 
 def test_setup_linear() -> None:
@@ -57,3 +57,16 @@ def test_train_linear_basic() -> None:
     # Expect the minimum goodness of the positive data to be bigger than the max goodness of neg.
     assert g_pos.min().item() > g_neg.max().item()
     assert g_pos.mean().item() > g_neg.mean().item()
+
+
+def test_sparsity_linear() -> None:
+    torch.manual_seed(42)
+    linear = FFLinear(in_features=10, out_features=2, loss_threshold=1, lr=0.02)
+    hoyer = linear.sparsity(SparsityType.HOYER).item()
+    entropy = linear.sparsity(SparsityType.ENTROPY_BASED).item()
+
+    print(f"Hoyer: {hoyer}")
+    print(f"Entropy-based: {entropy}")
+
+    assert hoyer >= 0 and hoyer <= 1
+    assert entropy >= 0 and entropy <= 1