X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?p=pytorch.git;a=blobdiff_plain;f=ae_size.py;h=067a7fa7bc1d44f4330e7f37fe1e511d3fd92045;hp=71839378d3001d1c5b5eae1601d64b3509712995;hb=HEAD;hpb=0a30c677bfb0aac1f9eebc9646728519ca72bf5b

diff --git a/ae_size.py b/ae_size.py
index 7183937..49f4a20 100755
--- a/ae_size.py
+++ b/ae_size.py
@@ -1,35 +1,42 @@
 #!/usr/bin/env python
 
+# Any copyright is dedicated to the Public Domain.
+# https://creativecommons.org/publicdomain/zero/1.0/
+
+# Written by Francois Fleuret <francois@fleuret.org>
+
 import math
 from torch import nn
 from torch import Tensor
 
 ######################################################################
 
+
 def minimal_input_size(w, layer_specs):
-    assert w > 0, 'The input is too small'
+    assert w > 0, "The input is too small"
     if layer_specs == []:
         return w
     else:
-        k, s = layer_specs[0]
-        w = math.ceil((w - k) / s) + 1
-        w = minimal_input_size(w, layer_specs[1:])
-        return int((w - 1) * s + k)
+        kernel_size, stride = layer_specs[0]
+        v = int(math.ceil((w - kernel_size) / stride)) + 1
+        v = minimal_input_size(v, layer_specs[1:])
+        return (v - 1) * stride + kernel_size
+
 
 ######################################################################
 
 # Dummy test
 
 if __name__ == "__main__":
-
-    layer_specs = [ (11, 5), (5, 2), (3, 2), (3, 2) ]
+    layer_specs = [(17, 5), (5, 4), (3, 2), (3, 2)]
 
     layers = []
-    for l in layer_specs:
-        layers.append(nn.Conv2d(1, 1, l[0], l[1]))
 
-    for l in reversed(layer_specs):
-        layers.append(nn.ConvTranspose2d(1, 1, l[0], l[1]))
+    for kernel_size, stride in layer_specs:
+        layers.append(nn.Conv2d(1, 1, kernel_size, stride))
+
+    for kernel_size, stride in reversed(layer_specs):
+        layers.append(nn.ConvTranspose2d(1, 1, kernel_size, stride))
 
     m = nn.Sequential(*layers)