fix kfac with bias

nngeometry/object/pspace.py

@@ -487,7 +487,7 @@ def get_dense_tensor(self, split_weight_bias=True):
             a, g = self.data[layer_id]
             start = self.generator.layer_collection.p_pos[layer_id]
             sAG = a.size(0) * g.size(0)
-            if split_weight_bias:
+            if split_weight_bias and layer.bias:
                 reconstruct = torch.cat(
                     [
                         torch.cat(
@@ -517,7 +517,7 @@ def get_diag(self, split_weight_bias=True):
         for layer_id, layer in self.generator.layer_collection.layers.items():
             a, g = self.data[layer_id]
             diag_of_block = torch.diag(g).view(-1, 1) * torch.diag(a).view(1, -1)
-            if split_weight_bias:
+            if split_weight_bias and layer.bias:
                 diags.append(diag_of_block[:, :-1].contiguous().view(-1))
                 diags.append(diag_of_block[:, -1:].view(-1))
             else:
@@ -535,10 +535,10 @@ def mv(self, vs):
                 v = torch.cat([v, vs_dict[layer_id][1].unsqueeze(1)], dim=1)
             a, g = self.data[layer_id]
             mv = torch.mm(torch.mm(g, v), a)
-            if layer.bias is None:
-                mv_tuple = (mv.view(*sw),)
-            else:
+            if layer.bias:
                 mv_tuple = (mv[:, :-1].contiguous().view(*sw), mv[:, -1].contiguous())
+            else:
+                mv_tuple = (mv.view(*sw),)
             out_dict[layer_id] = mv_tuple
         return PVector(layer_collection=vs.layer_collection, dict_repr=out_dict)
 

tests/tasks.py

@@ -101,7 +101,7 @@ def __init__(self, normalization="none"):
             self.bn1 = nn.BatchNorm2d(6)
         elif self.normalization == "group_norm":
             self.gn1 = nn.GroupNorm(2, 6)
-        self.conv2 = nn.Conv2d(6, 5, 4, 1)
+        self.conv2 = nn.Conv2d(6, 5, 4, 1, bias=False)
         self.conv3 = nn.Conv2d(5, 7, 3, 1, 1)
         if self.normalization == "weight_norm":
             self.wn2 = WeightNorm1d(7, 4)

tests/test_jacobian_kfac.py

@@ -4,7 +4,13 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from tasks import get_conv_task, get_fullyconnect_task, get_mnist, to_device_model
+from tasks import (
+    get_conv_task,
+    get_fullyconnect_task,
+    get_mnist,
+    to_device_model,
+    get_conv1d_task,
+)
 from torch.utils.data import DataLoader, Subset
 from utils import angle, check_ratio, check_tensors
 
@@ -161,7 +167,7 @@ def test_jacobian_kfac_vs_pblockdiag():
 
 
 def test_jacobian_kfac():
-    for get_task in [get_fullyconnect_task, get_conv_task]:
+    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
         loader, lc, parameters, model, function, n_output = get_task()
 
         generator = Jacobian(
@@ -190,6 +196,8 @@ def test_jacobian_kfac():
         # Test mv
         mv_direct = torch.mv(G_kfac_split, random_v.get_flat_representation())
         mv_kfac = M_kfac.mv(random_v)
+        print(mv_direct.size(), lc.layers)
+        pvec = PVector(layer_collection=lc, vector_repr=mv_direct)
         check_tensors(mv_direct, mv_kfac.get_flat_representation())
 
         # Test vTMv