3077874c1ea50310452bf15dd6b2f496a705eef6
[pytorch.git] / conv_chain.py
1 #!/usr/bin/env python
2
3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
5
6 # Written by Francois Fleuret <francois@fleuret.org>
7
8 ######################################################################
9
10 def conv_chain(input_size, output_size, remain_depth, cond):
11     if remain_depth == 0:
12         if input_size == output_size:
13             return [ [ ] ]
14         else:
15             return [ ]
16     else:
17         r = [ ]
18         for kernel_size in range(1, input_size + 1):
19             for stride in range(1, input_size):
20                 if cond(remain_depth, kernel_size, stride):
21                     n = (input_size - kernel_size) // stride + 1
22                     if n >= output_size and (n - 1) * stride + kernel_size == input_size:
23                         q = conv_chain(n, output_size, remain_depth - 1, cond)
24                         r += [ [ (kernel_size, stride) ] + u for u in q ]
25         return r
26
27 ######################################################################
28
29 if __name__ == "__main__":
30
31     import torch
32     from torch import nn
33
34     # Example
35
36     c = conv_chain(
37         input_size = 64, output_size = 8,
38         remain_depth = 5,
39         # We want kernels smaller than 4, strides smaller than the
40         # kernels, and strides of 1 except in the two last layers
41         cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2)
42     )
43
44     x = torch.rand(1, 1, 64)
45
46     for m in c:
47         model = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ])
48         print(model)
49         print(x.size(), model(x).size())
50
51 ######################################################################