extra str info in SRigL

mklasby · Aug 9, 2024 · 9f79587 · 9f79587
1 parent 88682ee
commit 9f79587
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/sparsimony/dst/base.py b/sparsimony/dst/base.py
@@ -279,7 +279,7 @@ def neuron_is_active(neuron):
             f"Global Sparsity Actual: {global_sparsity}\n"
             f"Layerwise Sparsity Targets: {layerwise_sparsity_target}\n"
             f"Layerwise Sparsity Actual: {layerwise_sparsity_actual}\n"
-            f"Active/Total Neurons: {active_vs_total_neurons}"
+            f"Active/Total Neurons: {active_vs_total_neurons}\n"
         )
 
     def __repr__(self) -> str:

diff --git a/sparsimony/dst/srigl.py b/sparsimony/dst/srigl.py
@@ -1,4 +1,6 @@
 from typing import Optional, Dict, Any
+from math import prod
+
 import torch
 import torch.nn as nn
 from torch.ao.pruning.sparsifier.base_sparsifier import BaseSparsifier
@@ -183,3 +185,13 @@ def _assert_sparsity_level(self, mask: torch.Tensor, sparsity_level: float):
                 "FFI error: Multiple non-zero FFI values detected: "
                 f"{ffi.unique()}"
             )
+
+    def __str__(self) -> str:
+        ffi = []
+        for config in self.groups:
+            mask = get_mask(**config)
+            mask_flat = mask.view(mask.shape[0], prod(mask.shape[1:]))
+            ffi.append(mask_flat.sum(dim=1, dtype=torch.int).unique().item())
+        s = super().__str__()
+        s += f"FFI: {ffi}\n"
+        return s