updated tests and started extensive testing

Author: thib-s

orthogonium/layers/channel_shuffle.py

@@ -24,20 +24,3 @@ def forward(self, x):
 
     def extra_repr(self):
         return f"group_in={self.group_in}, group_out={self.group_out}"
-
-
-if __name__ == "__main__":
-    x = torch.randn(2, 6)
-    # takes two groups of size 3, and return 3 groups of size 2
-    gm = ChannelShuffle(2, 3)
-    print(f"in: {x}")
-    y = gm(x)
-    print(f"out: {y}")
-    x2 = torch.randn(2, 6, 32, 32)
-    gm = ChannelShuffle(2, 3)
-    y2 = gm(x2)
-    print(f"in shape: {x2.shape}, out shape: {y2.shape}")
-    gp = ChannelShuffle(3, 2)
-    x2b = gp(y2)
-    assert torch.allclose(x2, x2b), "ChannelShuffle is not invertible"
-    print("ChannelShuffle is invertible")

orthogonium/layers/conv/AOC/bcop_x_rko_conv.py

@@ -54,24 +54,20 @@ def __init__(
         if (self.dilation[0] != 1 or self.dilation[1] != 1) and (
             self.stride[0] != 1 or self.stride[1] != 1
         ):
-            raise RuntimeError(
+            raise ValueError(
                 "dilation must be 1 when stride is not 1. The set of orthonal convolutions is empty in this setting."
             )
         # raise runtime error if kernel size >= stride
         if self.kernel_size[0] < self.stride[0] or self.kernel_size[1] < self.stride[1]:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
-        if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
-            )
         if (
             ((max(in_channels, out_channels) // groups) < 2)
             and (self.kernel_size[0] != self.stride[0])
             and (self.kernel_size[1] != self.stride[1])
         ):
-            raise RuntimeError("inner conv must have at least 2 channels")
+            raise ValueError("inner conv must have at least 2 channels")
         self.intermediate_channels = max(
             in_channels, out_channels // (self.stride[0] * self.stride[1])
         )
@@ -207,23 +203,19 @@ def __init__(
         if (self.dilation[0] != 1 or self.dilation[1] != 1) and (
             self.stride[0] != 1 or self.stride[1] != 1
         ):
-            raise RuntimeError(
+            raise ValueError(
                 "dilation must be 1 when stride is not 1. The set of orthonal convolutions is empty in this setting."
             )
         if self.kernel_size[0] < self.stride[0] or self.kernel_size[1] < self.stride[1]:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
-        if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
-            )
         if (
             ((max(in_channels, out_channels) // groups) < 2)
             and (self.kernel_size[0] != self.stride[0])
             and (self.kernel_size[1] != self.stride[1])
         ):
-            raise RuntimeError("inner conv must have at least 2 channels")
+            raise ValueError("inner conv must have at least 2 channels")
         if out_channels * (self.stride[0] * self.stride[1]) >= in_channels:
             self.intermediate_channels = max(
                 in_channels // (self.stride[0] * self.stride[1]), out_channels

orthogonium/layers/conv/AOC/fast_block_ortho_conv.py

@@ -390,22 +390,24 @@ def __init__(
         )
 
         # raise runtime error if kernel size >= stride
+        if self.kernel_size[0] < self.stride[0] or self.kernel_size[1] < self.stride[1]:
+            raise ValueError(
+                "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
+            )
         if (
             (self.stride[0] > 1 or self.stride[1] > 1) and (out_channels > in_channels)
         ) or (
             self.stride[0] > self.kernel_size[0] or self.stride[1] > self.kernel_size[1]
         ):
-            raise RuntimeError(
+            raise ValueError(
                 "stride > 1 is not supported when out_channels > in_channels, "
                 "use TODO layer instead"
             )
-        if self.kernel_size[0] < self.stride[0] or self.kernel_size[1] < self.stride[1]:
-            raise RuntimeError(
-                "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
-            )
-        if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
+        if (self.dilation[0] != 1 or self.dilation[1] != 1) and (
+            self.stride[0] != 1 or self.stride[1] != 1
+        ):
+            raise ValueError(
+                "dilation must be 1 when stride is not 1. The set of orthonal convolutions is empty in this setting."
             )
         del self.weight
         attach_bcop_weight(
@@ -492,19 +494,15 @@ def __init__(
         )
         # raise runtime error if kernel size >= stride
         if self.kernel_size[0] < self.stride[0] or self.kernel_size[1] < self.stride[1]:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
-        if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
-            )
         if (
             ((max(in_channels, out_channels) // groups) < 2)
             and (self.kernel_size[0] != self.stride[0])
             and (self.kernel_size[1] != self.stride[1])
         ):
-            raise RuntimeError("inner conv must have at least 2 channels")
+            raise ValueError("inner conv must have at least 2 channels")
         if out_channels * (self.stride[0] * self.stride[1]) < in_channels:
             # raise warning because this configuration don't yield orthogonal
             # convolutions
@@ -514,6 +512,10 @@ def __init__(
                 "transposed convolutions",
                 RuntimeWarning,
             )
+        if (self.dilation[0] > 1 or self.dilation[1] > 1) and (
+            self.stride[0] != 1 or self.stride[1] != 1
+        ):
+            raise ValueError("dilation must be 1 when stride is not 1")
         del self.weight
         attach_bcop_weight(
             self,

orthogonium/layers/conv/AOC/ortho_conv.py

@@ -34,7 +34,7 @@ def AdaptiveOrthoConv2d(
 
     """
     if kernel_size < stride:
-        raise RuntimeError(
+        raise ValueError(
             "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
         )
     if kernel_size == stride:
@@ -83,7 +83,7 @@ def AdaptiveOrthoConvTranspose2d(
         `out_channels * (stride**2) > in_channels`.
     """
     if kernel_size < stride:
-        raise RuntimeError(
+        raise ValueError(
             "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
         )
     if kernel_size == stride:

orthogonium/layers/conv/AOC/rko_conv.py

@@ -54,7 +54,9 @@ def forward(self, X):
         )
         X = self.pi(X)
         X = self.bjorck(X)
-        X = X.view(self.out_channels, self.in_channels // self.groups, self.k1, self.k2)
+        X = X.reshape(
+            self.out_channels, self.in_channels // self.groups, self.k1, self.k2
+        )
         return X * self.scale
 
     def right_inverse(self, X):
@@ -120,11 +122,10 @@ def __init__(
             bias,
             padding_mode,
         )
-        if self.dilation[0] > 1 or self.dilation[1] > 1:
-            raise RuntimeWarning(
-                "Dilation must be 1 in the RKO convolution."
-                "Use RkoConvTranspose2d instead."
-            )
+        if (self.dilation[0] > 1 or self.dilation[1] > 1) and (
+            self.stride[0] != 1 or self.stride[1] != 1
+        ):
+            raise ValueError("dilation must be 1 when stride is not 1")
         # torch.nn.init.orthogonal_(self.weight)
         self.scale = 1 / math.sqrt(
             math.ceil(self.dilation[0] * self.kernel_size[0] / self.stride[0])
@@ -172,7 +173,7 @@ def singular_values(self):
             return sv_min, sv_max, stable_rank
         elif self.stride[0] > 1 or self.stride[1] > 1:
             raise RuntimeError(
-                "Not able to compute singular values for this " "configuration"
+                "Not able to compute singular values for this configuration"
             )
         # Implements interface required by LipschitzModuleL2
         sv_min, sv_max, stable_rank = conv_singular_values_numpy(
@@ -220,15 +221,14 @@ def __init__(
         )
 
         # raise runtime error if kernel size >= stride
-        if self.kernel_size[0] > self.stride[0] or self.kernel_size[1] > self.stride[1]:
-            raise RuntimeError(
-                "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
+        if self.kernel_size[0] < self.stride[0] or self.kernel_size[1] < self.stride[1]:
+            raise ValueError(
+                "kernel size must be smaller than stride. The set of orthogonal convolutions is empty in this setting."
             )
-        if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
-            )
-
+        if (self.dilation[0] > 1 or self.dilation[1] > 1) and (
+            self.stride[0] != 1 or self.stride[1] != 1
+        ):
+            raise ValueError("dilation must be 1 when stride is not 1")
         if (
             self.stride[0] != self.kernel_size[0]
             or self.stride[1] != self.kernel_size[1]
@@ -271,7 +271,7 @@ def singular_values(self):
             return sv_min, sv_max, stable_rank
         elif self.stride[0] > 1 or self.stride[1] > 1:
             raise RuntimeError(
-                "Not able to compute singular values for this " "configuration"
+                "Not able to compute singular values for this configuration"
             )
         # Implements interface required by LipschitzModuleL2
         sv_min, sv_max, stable_rank = conv_singular_values_numpy(

orthogonium/layers/conv/SLL/sll_layer.py

@@ -200,37 +200,32 @@ def __init__(
     def forward(self, x):
         kernel = self.in_conv.weight
         # conv
+        res = x
         if self.padding_mode == "circular":
-            x = F.pad(
-                x,
-                (
-                    self.kernel_size // 2,
-                    self.kernel_size // 2,
-                    self.kernel_size // 2,
-                    self.kernel_size // 2,
-                ),
+            res = F.pad(
+                res,
+                (self.padding,) * 4,
                 mode="circular",
+                value=0,
             )
         res = F.conv2d(
-            x, kernel, bias=self.in_conv.bias, padding=self.padding, groups=self.groups
+            res,
+            kernel,
+            bias=self.in_conv.bias,
+            padding=0,
+            groups=self.groups,
         )
         # activation
         res = self.activation(res)
         # conv transpose
         if self.padding_mode == "circular":
             res = F.pad(
                 res,
-                (
-                    self.kernel_size // 2,
-                    self.kernel_size // 2,
-                    self.kernel_size // 2,
-                    self.kernel_size // 2,
-                ),
+                (self.padding,) * 4,
                 mode="circular",
+                value=0,
             )
-        res = 2 * F.conv_transpose2d(
-            res, kernel, padding=self.padding, groups=self.groups
-        )
+        res = 2 * F.conv_transpose2d(res, kernel, padding=0, groups=self.groups)
         # residual
         out = x - res
         return out

orthogonium/layers/conv/adaptiveSOC/fast_skew_ortho_conv.py

@@ -230,17 +230,15 @@ def __init__(
 
         # raise runtime error if kernel size >= stride
         if ((stride > 1) and (out_channels > in_channels)) or (stride > kernel_size):
-            raise RuntimeError(
+            raise ValueError(
                 "stride > 1 is not supported when out_channels > in_channels, "
                 "use TODO layer instead"
             )
         if kernel_size < stride:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
         if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
             )
         self.padding = padding
         self.stride = stride
@@ -340,15 +338,13 @@ def __init__(
 
         # raise runtime error if kernel size >= stride
         if kernel_size < stride:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
         if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
             )
         if ((self.max_channels // groups) < 2) and (kernel_size != stride):
-            raise RuntimeError("inner conv must have at least 2 channels")
+            raise ValueError("inner conv must have at least 2 channels")
         if out_channels * (stride**2) < in_channels:
             # raise warning because this configuration don't yield orthogonal
             # convolutions

orthogonium/layers/conv/adaptiveSOC/ortho_conv.py

@@ -35,7 +35,7 @@ def AdaptiveSOCConv2d(
     Otherwise, the layer is a BcopRkoConv2d.
     """
     if kernel_size < stride:
-        raise RuntimeError(
+        raise ValueError(
             "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
         )
     if kernel_size == stride:
@@ -84,7 +84,7 @@ def AdaptiveSOCConvTranspose2d(
         `out_channels * (stride**2) > in_channels`.
     """
     if kernel_size < stride:
-        raise RuntimeError(
+        raise ValueError(
             "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
         )
     if kernel_size == stride:

orthogonium/layers/conv/adaptiveSOC/soc_x_rko_conv.py

@@ -64,15 +64,13 @@ def __init__(
 
         # raise runtime error if kernel size >= stride
         if kernel_size < stride:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
         if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
             )
         if ((self.max_channels // groups) < 2) and (kernel_size != stride):
-            raise RuntimeError("inner conv must have at least 2 channels")
+            raise ValueError("inner conv must have at least 2 channels")
         self.padding = padding
         self.stride = stride
         self.kernel_size = kernel_size
@@ -194,8 +192,6 @@ def __init__(
         But unit testing have shown that the convolution is still orthogonal when
          `out_channels * (stride**2) > in_channels`.
         """
-        if dilation != 1:
-            raise RuntimeError("dilation not supported")
         super(SOCRkoConvTranspose2d, self).__init__(
             in_channels,
             out_channels,
@@ -219,15 +215,13 @@ def __init__(
 
         # raise runtime error if kernel size >= stride
         if kernel_size < stride:
-            raise RuntimeError(
+            raise ValueError(
                 "kernel size must be smaller than stride. The set of orthonal convolutions is empty in this setting."
             )
         if (in_channels % groups != 0) and (out_channels % groups != 0):
-            raise RuntimeError(
-                "in_channels and out_channels must be divisible by groups"
             )
         if ((self.max_channels // groups) < 2) and (kernel_size != stride):
-            raise RuntimeError("inner conv must have at least 2 channels")
+            raise ValueError("inner conv must have at least 2 channels")
         self.padding = padding
         self.stride = stride
         self.kernel_size = kernel_size