X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=conv_chain.py;h=2d5af8ecf616f1e2b9c37e206c16998e40cc58a5;hb=2d95f238bbaa0e585b50846d39c98df4aae2b7f9;hp=85ae4fcc5d1040f985538791af35cb8631ad29e4;hpb=77af470cdb98b4a97af6432ca1421c28062b9aae;p=pytorch.git

diff --git a/conv_chain.py b/conv_chain.py
index 85ae4fc..2d5af8e 100755
--- a/conv_chain.py
+++ b/conv_chain.py
@@ -14,10 +14,10 @@ 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):
+            for stride in range(1, input_size):
                 if cond(depth, kernel_size, stride):
                     n = (input_size - kernel_size) // stride + 1
-                    if (n - 1) * stride + kernel_size == input_size:
+                    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
@@ -31,7 +31,9 @@ if __name__ == "__main__":
     c = conv_chain(
         input_size = 64, output_size = 8,
         depth = 5,
-        cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d < 3)
+        # We want kernels smaller than 4, strides smaller than the
+        # kernels, and stride 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)