#!/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>
######################################################################
-def conv_chain(input_size, output_size, depth, cond):
- if depth == 0:
+def conv_chain(input_size, output_size, remain_depth, cond):
+ if remain_depth == 0:
if input_size == output_size:
return [ [ ] ]
else:
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(remain_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, remain_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,
+ remain_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())
######################################################################