From: Francois Fleuret <francois@fleuret.org>
Date: Sun, 6 Jun 2021 12:52:34 +0000 (+0200)
Subject: Update.
X-Git-Url: https://www.fleuret.org/cgi-bin/gitweb/gitweb.cgi?p=pytorch.git;a=commitdiff_plain;h=77af470cdb98b4a97af6432ca1421c28062b9aae

Update.
---

diff --git a/conv_chain.py b/conv_chain.py
index fa5d752..85ae4fc 100755
--- a/conv_chain.py
+++ b/conv_chain.py
@@ -15,10 +15,10 @@ def conv_chain(input_size, output_size, depth, cond):
         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)
+                if cond(depth, kernel_size, stride):
+                    n = (input_size - kernel_size) // stride + 1
+                    if (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,7 @@ if __name__ == "__main__":
     c = conv_chain(
         input_size = 64, output_size = 8,
         depth = 5,
-        cond = lambda k, s: k <= 4 and s <= 2 and s <= k//2
+        cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d < 3)
     )
 
     x = torch.rand(1, 1, 64)