adds FIM helper for regression (a.k.a. Gauss Newton)

tfjgeorge · Jul 29, 2020 · 815bb1c · 815bb1c
1 parent 280ba08
commit 815bb1c
Showing 1 changed file with 13 additions and 6 deletions.
diff --git a/nngeometry/metrics.py b/nngeometry/metrics.py
@@ -53,11 +53,18 @@ def function(*d):
             probs = softmax(logits, dim=1).detach()
             return (logits * probs**.5 * (1 - probs))
 
-        generator = Jacobian(layer_collection=layer_collection,
-                             model=model,
-                             loader=loader,
-                             function=function,
-                             n_output=n_output)
-        return representation(generator)
+    elif variant == 'regression':
+
+        def function(*d):
+            inputs = d[0].to(device)
+            estimates = model(inputs)
+            return estimates
     else:
         raise NotImplementedError
+
+    generator = Jacobian(layer_collection=layer_collection,
+                         model=model,
+                         loader=loader,
+                         function=function,
+                         n_output=n_output)
+    return representation(generator)