Merge pull request #76 from tfjgeorge/conv1d

Conv1d

nngeometry/generator/jacobian/grads.py

@@ -1,13 +1,21 @@
 import torch
 import torch.nn.functional as F
 
-from nngeometry.layercollection import (Affine1dLayer, BatchNorm1dLayer,
-                                        BatchNorm2dLayer, Conv2dLayer,
-                                        ConvTranspose2dLayer, Cosine1dLayer,
-                                        GroupNormLayer, LinearLayer,
-                                        WeightNorm1dLayer, WeightNorm2dLayer)
-
-from .grads_conv import conv2d_backward, convtranspose2d_backward
+from nngeometry.layercollection import (
+    Affine1dLayer,
+    BatchNorm1dLayer,
+    BatchNorm2dLayer,
+    Conv2dLayer,
+    ConvTranspose2dLayer,
+    Cosine1dLayer,
+    GroupNormLayer,
+    LinearLayer,
+    WeightNorm1dLayer,
+    WeightNorm2dLayer,
+    Conv1dLayer,
+)
+
+from .grads_conv import conv2d_backward, convtranspose2d_backward, conv1d_backward
 
 
 class JacobianFactory:
@@ -325,8 +333,90 @@ def flat_grad(cls, buffer, mod, layer, x, gy):
             buffer[:, w_numel:].add_(gy)
 
 
+class Conv1dJacobianFactory(JacobianFactory):
+    @classmethod
+    def flat_grad(cls, buffer, mod, layer, x, gy):
+        bs = x.size(0)
+        w_numel = layer.weight.numel()
+        indiv_gw = conv1d_backward(mod, x, gy)
+        buffer[:, :w_numel].add_(indiv_gw.view(bs, -1))
+        if layer.bias is not None:
+            buffer[:, w_numel:].add_(gy.sum(dim=2))
+
+    @classmethod
+    def Jv(cls, buffer, mod, layer, x, gy, v, v_bias):
+        bs = x.size(0)
+        gy2 = F.conv1d(
+            x, v, stride=mod.stride, padding=mod.padding, dilation=mod.dilation
+        )
+        buffer.add_((gy * gy2).view(bs, -1).sum(dim=1))
+        if layer.bias is not None:
+            buffer.add_(torch.mv(gy.sum(dim=2), v_bias))
+
+    @classmethod
+    def kfac_xx(cls, buffer, mod, layer, x, gy):
+        ks = (1, mod.weight.size(2))
+        # A_tilda in KFC
+        A_tilda = F.unfold(
+            x.unsqueeze(2),
+            kernel_size=ks,
+            stride=(1, mod.stride[0]),
+            padding=(0, mod.padding[0]),
+            dilation=(1, mod.dilation[0]),
+        )
+        # A_tilda is bs * #locations x #parameters
+        A_tilda = A_tilda.permute(0, 2, 1).contiguous().view(-1, A_tilda.size(1))
+        if layer.bias is not None:
+            A_tilda = torch.cat([A_tilda, torch.ones_like(A_tilda[:, :1])], dim=1)
+        # Omega_hat in KFC
+        buffer.add_(torch.mm(A_tilda.t(), A_tilda))
+
+    @classmethod
+    def kfac_gg(cls, buffer, mod, layer, x, gy):
+        spatial_locations = gy.size(2)
+        os = gy.size(1)
+        # DS_tilda in KFC
+        DS_tilda = gy.permute(0, 2, 1).contiguous().view(-1, os)
+        buffer.add_(torch.mm(DS_tilda.t(), DS_tilda) / spatial_locations)
+
+    @classmethod
+    def kfe_diag(cls, buffer, mod, layer, x, gy, evecs_a, evecs_g):
+        ks = (1, mod.weight.size(2))
+        gy_s = gy.size()
+        bs = gy_s[0]
+        # project x to kfe
+        x_unfold = F.unfold(
+            x.unsqueeze(2),
+            kernel_size=ks,
+            stride=(1, mod.stride[0]),
+            padding=(0, mod.padding[0]),
+            dilation=(1, mod.dilation[0]),
+        )
+        x_unfold_s = x_unfold.size()
+        x_unfold = (
+            x_unfold.view(bs, x_unfold_s[1], -1)
+            .permute(0, 2, 1)
+            .contiguous()
+            .view(-1, x_unfold_s[1])
+        )
+        if mod.bias is not None:
+            x_unfold = torch.cat([x_unfold, torch.ones_like(x_unfold[:, :1])], dim=1)
+        x_kfe = torch.mm(x_unfold, evecs_a)
+
+        # project gy to kfe
+        gy = gy.view(bs, gy_s[1], -1).permute(0, 2, 1).contiguous()
+        gy_kfe = torch.mm(gy.view(-1, gy_s[1]), evecs_g)
+        gy_kfe = gy_kfe.view(bs, -1, gy_s[1]).permute(0, 2, 1).contiguous()
+
+        indiv_gw = torch.bmm(
+            gy_kfe.view(bs, gy_s[1], -1), x_kfe.view(bs, -1, x_kfe.size(1))
+        )
+        buffer.add_((indiv_gw**2).sum(dim=0).view(-1))
+
+
 FactoryMap = {
     LinearLayer: LinearJacobianFactory,
+    Conv1dLayer: Conv1dJacobianFactory,
     Conv2dLayer: Conv2dJacobianFactory,
     ConvTranspose2dLayer: ConvTranspose2dJacobianFactory,
     BatchNorm1dLayer: BatchNorm1dJacobianFactory,

nngeometry/generator/jacobian/grads_conv.py

@@ -83,11 +83,6 @@ def conv_backward(
     return weight_bgrad
 
 
-def conv1d_backward(*args, **kwargs):
-    """Computes per-example gradients for nn.Conv1d layers."""
-    return conv_backward(*args, nd=1, **kwargs)
-
-
 def conv2d_backward_using_conv(mod, x, gy):
     """Computes per-example gradients for nn.Conv2d layers."""
     return conv_backward(
@@ -119,7 +114,29 @@ def conv2d_backward_using_unfold(mod, x, gy):
 
 
 def conv2d_backward(*args, **kwargs):
-    return _conv_grad_impl.get_impl()(*args, **kwargs)
+    return _conv_grad_impl.get_impl2d()(*args, **kwargs)
+
+
+def conv1d_backward_using_unfold(mod, x, gy):
+    """Computes per-example gradients for nn.Conv1d layers."""
+    ks = (1, mod.weight.size(2))
+    gy_s = gy.size()
+    bs = gy_s[0]
+    x_unfold = F.unfold(
+        x.unsqueeze(2),
+        kernel_size=ks,
+        stride=(1, mod.stride[0]),
+        padding=(0, mod.padding[0]),
+        dilation=(1, mod.dilation[0]),
+    )
+    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 conv1d_backward(*args, **kwargs):
+    return _conv_grad_impl.get_impl1d()(*args, **kwargs)
 
 
 class ConvGradImplManager:
@@ -129,12 +146,18 @@ def __init__(self):
     def use_unfold(self, choice=True):
         self._use_unfold = choice
 
-    def get_impl(self):
+    def get_impl2d(self):
         if self._use_unfold:
             return conv2d_backward_using_unfold
         else:
             return conv2d_backward_using_conv
 
+    def get_impl1d(self):
+        if self._use_unfold:
+            return conv1d_backward_using_unfold
+        else:
+            raise NotImplementedError()
+
 
 _conv_grad_impl = ConvGradImplManager()
 

nngeometry/generator/jacobian/jacobian.py

@@ -41,7 +41,10 @@ def __init__(
         self.centering = centering
 
         if function is None:
-            function = lambda *x: model(x[0])
+
+            def function(*x):
+                return model(x[0])
+
         self.function = function
 
         if layer_collection is None:
@@ -250,6 +253,9 @@ def get_kfac_blocks(self, examples):
             elif layer_class == "Conv2dLayer":
                 sG = layer.out_channels
                 sA = layer.in_channels * layer.kernel_size[0] * layer.kernel_size[1]
+            elif layer_class == "Conv1dLayer":
+                sG = layer.out_channels
+                sA = layer.in_channels * layer.kernel_size[0]
             if layer.bias is not None:
                 sA += 1
             self._blocks[layer_id] = (
@@ -459,6 +465,9 @@ def get_kfe_diag(self, kfe, examples):
             elif layer_class == "Conv2dLayer":
                 sG = layer.out_channels
                 sA = layer.in_channels * layer.kernel_size[0] * layer.kernel_size[1]
+            elif layer_class == "Conv1dLayer":
+                sG = layer.out_channels
+                sA = layer.in_channels * layer.kernel_size[0]
             if layer.bias is not None:
                 sA += 1
             self._diags[layer_id] = torch.zeros((sG * sA), device=device, dtype=dtype)
@@ -761,7 +770,7 @@ def _hook_compute_kfac_blocks(self, mod, gy):
         layer_id = self.m_to_l[mod]
         layer = self.layer_collection[layer_id]
         block = self._blocks[layer_id]
-        if mod_class in ["Linear", "Conv2d"]:
+        if mod_class in ["Linear", "Conv2d", "Conv1d"]:
             FactoryMap[layer.__class__].kfac_gg(block[1], mod, layer, x, gy)
             if self.i_output == 0:
                 # do this only once if n_output > 1
@@ -775,7 +784,7 @@ def _hook_compute_kfe_diag(self, mod, gy):
         layer = self.layer_collection[layer_id]
         x = self.xs[mod]
         evecs_a, evecs_g = self._kfe[layer_id]
-        if mod_class in ["Linear", "Conv2d"]:
+        if mod_class in ["Linear", "Conv2d", "Conv1d"]:
             FactoryMap[layer.__class__].kfe_diag(
                 self._diags[layer_id], mod, layer, x, gy, evecs_a, evecs_g
             )

nngeometry/layercollection.py

@@ -24,6 +24,7 @@ class LayerCollection:
         "Cosine1d",
         "Affine1d",
         "ConvTranspose2d",
+        "Conv1d",
     ]
 
     def __init__(self, layers=None):
@@ -112,6 +113,13 @@ def _module_to_layer(mod):
                 kernel_size=mod.kernel_size,
                 bias=(mod.bias is not None),
             )
+        elif mod_class == "Conv1d":
+            return Conv1dLayer(
+                in_channels=mod.in_channels,
+                out_channels=mod.out_channels,
+                kernel_size=mod.kernel_size,
+                bias=(mod.bias is not None),
+            )
         elif mod_class == "BatchNorm1d":
             return BatchNorm1dLayer(num_features=mod.num_features)
         elif mod_class == "BatchNorm2d":
@@ -231,6 +239,31 @@ def __eq__(self, other):
         )
 
 
+class Conv1dLayer(AbstractLayer):
+    def __init__(self, in_channels, out_channels, kernel_size, bias=True):
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.weight = Parameter(out_channels, in_channels, kernel_size[0])
+        if bias:
+            self.bias = Parameter(out_channels)
+        else:
+            self.bias = None
+
+    def numel(self):
+        if self.bias is not None:
+            return self.weight.numel() + self.bias.numel()
+        else:
+            return self.weight.numel()
+
+    def __eq__(self, other):
+        return (
+            self.in_channels == other.in_channels
+            and self.out_channels == other.out_channels
+            and self.kernel_size == other.kernel_size
+        )
+
+
 class LinearLayer(AbstractLayer):
     def __init__(self, in_features, out_features, bias=True):
         self.in_features = in_features

nngeometry/object/__init__.py

@@ -1,7 +1,15 @@
 from .fspace import FMatDense
 from .map import PullBackDense, PushForwardDense, PushForwardImplicit
-from .pspace import (PMatBlockDiag, PMatDense, PMatDiag, PMatEKFAC,
-                     PMatImplicit, PMatKFAC, PMatLowRank, PMatQuasiDiag)
+from .pspace import (
+    PMatBlockDiag,
+    PMatDense,
+    PMatDiag,
+    PMatEKFAC,
+    PMatImplicit,
+    PMatKFAC,
+    PMatLowRank,
+    PMatQuasiDiag,
+)
 from .vector import FVector, PVector
 
 __all__ = [

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)
 
@@ -602,7 +602,7 @@ class PMatEKFAC(PMatAbstract):
     """
     EKFAC representation from
     *George, Laurent et al., Fast Approximate Natural Gradient Descent
-    in a Kronecker-factored Eigenbasis, NIPS 2018*
+    in a Kronecker-factored Eigenbasis, NeurIPS 2018*
 
     """
 
@@ -659,9 +659,9 @@ def get_KFE(self, split_weight_bias=True):
         """
         evecs, _ = self.data
         KFE = dict()
-        for layer_id, _ in self.generator.layer_collection.layers.items():
+        for layer_id, layer in self.generator.layer_collection.layers.items():
             evecs_a, evecs_g = evecs[layer_id]
-            if split_weight_bias:
+            if split_weight_bias and layer.bias:
                 kronecker(evecs_g, evecs_a[:-1, :])
                 kronecker(evecs_g, evecs_a[-1:, :].contiguous())
                 KFE[layer_id] = torch.cat(

nngeometry/object/vector.py

@@ -194,7 +194,7 @@ def _dict_to_flat(self):
         parts = []
         for layer_id, layer in self.layer_collection.layers.items():
             parts.append(self.dict_repr[layer_id][0].view(-1))
-            if len(self.dict_repr[layer_id]) > 1:
+            if layer.bias:
                 parts.append(self.dict_repr[layer_id][1].view(-1))
         return torch.cat(parts)
 

setup.py

@@ -11,5 +11,5 @@
                 'nngeometry.generator',
                 'nngeometry.generator.jacobian',
                 'nngeometry.object'],
-      install_requires=['torch>=2.0.0'],
+      install_requires=['torch>=2.0.0','torchvision>=0.9.1'],
       zip_safe=False)