Statistics computation for weight and bias matrices.

Author: mitkonikov

src/fflib/interfaces/iff_recurrent_layer.py

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

src/fflib/nn/ff_linear.py

@@ -3,9 +3,9 @@
 from torch.optim import Adam, Optimizer
 
 from fflib.enums import SparsityType
-from fflib.utils.maths import ComputeSparsity
+from fflib.utils.maths import ComputeSparsity, ComputeStats
 
-from typing import Callable, Tuple, Any, cast
+from typing import Callable, Tuple, Dict, Any, cast
 
 
 class FFLinear(Linear):
@@ -140,3 +140,6 @@ def strip_down(self) -> None:
     def sparsity(self, type: SparsityType) -> torch.Tensor:
         """Computes the sparsity of the weight's matrix."""
         return ComputeSparsity(torch.flatten(self.weight), type)
+
+    def stats(self) -> Dict[str, Any]:
+        return {"weight": ComputeStats(self.weight), "bias": ComputeStats(self.bias)}

src/fflib/nn/ff_net.py

@@ -118,3 +118,6 @@ def sparsity(self, type: SparsityType) -> Dict[str, float]:
         return {
             f"layer_{i}": float(layer.sparsity(type).item()) for i, layer in enumerate(self.layers)
         }
+
+    def stats(self) -> Dict[str, Any]:
+        return {f"layer_{i}": layer.stats() for i, layer in enumerate(self.layers)}

src/fflib/nn/ff_recurrent_layer.py

@@ -5,7 +5,7 @@
 from torch.optim import Adam, Optimizer
 from fflib.interfaces.iff_recurrent_layer import IFFRecurrentLayer
 from fflib.enums import SparsityType
-from fflib.utils.maths import ComputeSparsity
+from fflib.utils.maths import ComputeSparsity, ComputeStats
 from math import sqrt
 from typing import Callable, List, Tuple, Dict, cast, Any
 
@@ -164,6 +164,13 @@ def sparsity(self, type: SparsityType) -> Dict[str, float]:
             ),
         }
 
+    def stats(self) -> Dict[str, Any]:
+        return {
+            "fw": ComputeStats(self.fw),
+            "bw": ComputeStats(self.bw),
+            "fb": ComputeStats(self.fb),
+        }
+
 
 class FFRecurrentLayerDummy(IFFRecurrentLayer):
     def __init__(self, dimensions: int):
@@ -204,3 +211,6 @@ def strip_down(self) -> None:
 
     def sparsity(self, type: SparsityType) -> Dict[str, float]:
         return {}
+
+    def stats(self) -> Dict[str, Any]:
+        return {}

src/fflib/nn/ff_rnn.py

@@ -236,3 +236,7 @@ def strip_down(self) -> None:
     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)}
+
+    def stats(self) -> Dict[str, Dict[str, float]]:
+        """Returns a dictionary of dictionaries containing basic statistics about the layers' parameters."""
+        return {f"layer_{i}": layer.stats() for i, layer in enumerate(self.layers[1:-1])}

src/fflib/nn/ffc.py

@@ -5,7 +5,7 @@
 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 fflib.utils.maths import ComputeSparsity, ComputeStats
 
 from typing import List, Dict, Callable, Any
 
@@ -106,3 +106,11 @@ def sparsity(self, type: SparsityType) -> Dict[str, float]:
             ComputeSparsity(torch.flatten(self.classifier.weight), type).item()
         )
         return result
+
+    def stats(self) -> Dict[str, Any]:
+        result = {f"layer_{i}": layer.stats() for i, layer in enumerate(self.layers)}
+        result["classifier"] = {
+            "weight": ComputeStats(self.classifier.weight),
+            "bias": ComputeStats(self.classifier.bias),
+        }
+        return result

src/fflib/utils/maths.py

@@ -37,3 +37,12 @@ def ComputeAllSparsityTypes(x: torch.Tensor) -> Dict[str, torch.Tensor]:
     for type in SparsityType:
         result[str(type).split(".")[1]] = ComputeSparsity(x, type)
     return result
+
+
+def ComputeStats(x: torch.Tensor) -> Dict[str, float]:
+    return {
+        "min": x.min().item(),
+        "max": x.max().item(),
+        "mean": x.mean().item(),
+        "std": x.std().item(),
+    }