X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=conv_chain.py;h=d10798f4b2fd58b97b31283f4a9301fba210bff1;hb=0c2da9bc90d21bdbcbcbeeb0a070653cbc3c25cc;hp=04dfdfac0b32aaee16d81e9a46b6194303cbbb76;hpb=5ab8211805831629148d7b436b8770590f1987b0;p=pytorch.git

diff --git a/conv_chain.py b/conv_chain.py
index 04dfdfa..d10798f 100755
--- a/conv_chain.py
+++ b/conv_chain.py
@@ -1,7 +1,9 @@
 #!/usr/bin/env python
 
-import torch
-from torch import nn
+# Any copyright is dedicated to the Public Domain.
+# https://creativecommons.org/publicdomain/zero/1.0/
+
+# Written by Francois Fleuret <francois@fleuret.org>
 
 ######################################################################
 
@@ -14,29 +16,36 @@ def conv_chain(input_size, output_size, depth, cond):
     else:
         r = [ ]
         for kernel_size in range(1, input_size + 1):
-            for stride in range(1, input_size + 1):
-                if cond(kernel_size, stride):
-                    n = (input_size - kernel_size) // stride
-                    if n * stride + kernel_size == input_size:
-                        q = conv_chain(n + 1, output_size, depth - 1, cond)
+            for stride in range(1, input_size):
+                if cond(depth, kernel_size, stride):
+                    n = (input_size - kernel_size) // stride + 1
+                    if n >= output_size and (n - 1) * stride + kernel_size == input_size:
+                        q = conv_chain(n, output_size, depth - 1, cond)
                         r += [ [ (kernel_size, stride) ] + u for u in q ]
         return r
 
 ######################################################################
 
-# Example
+if __name__ == "__main__":
+
+    import torch
+    from torch import nn
+
+    # Example
 
-c = conv_chain(
-    input_size = 64, output_size = 8,
-    depth = 5,
-    cond = lambda k, s: k <= 4 and s <= 2 and s <= k//2
-)
+    c = conv_chain(
+        input_size = 64, output_size = 8,
+        depth = 5,
+        # We want kernels smaller than 4, strides smaller than the
+        # kernels, and strides of 1 except in the two last layers
+        cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2)
+    )
 
-x = torch.rand(1, 1, 64)
+    x = torch.rand(1, 1, 64)
 
-for m in c:
-    m = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ])
-    print(m)
-    print(x.size(), m(x).size())
+    for m in c:
+        model = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ])
+        print(model)
+        print(x.size(), model(x).size())
 
 ######################################################################