adds a function for displaying the dense matrix represented by any re…

…presentation
tfjgeorge · Aug 24, 2020 · 37ec416 · 37ec416
1 parent db8da63
commit 37ec416
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/nngeometry/utils.py b/nngeometry/utils.py
@@ -15,3 +15,13 @@ def per_example_grad_conv(mod, x, gy):
     x_unfold_s = x_unfold.size()
     return torch.bmm(gy.view(bs, gy_s[1], -1),
                      x_unfold.view(bs, x_unfold_s[1], -1).permute(0, 2, 1))
+
+
+def display_correl(M, axis):
+
+    M = M.get_dense_tensor()
+    diag = torch.diag(M)
+    dM = (diag + diag.mean() / 100) **.5
+    correl = torch.abs(M) / dM.unsqueeze(0) / dM.unsqueeze(1)
+
+    axis.imshow(correl.cpu())