85ae4fcc5d1040f985538791af35cb8631ad29e4
[pytorch.git] / conv_chain.py
1 #!/usr/bin/env python
2
3 import torch
4 from torch import nn
5
6 ######################################################################
7
8 def conv_chain(input_size, output_size, depth, cond):
9     if depth == 0:
10         if input_size == output_size:
11             return [ [ ] ]
12         else:
13             return [ ]
14     else:
15         r = [ ]
16         for kernel_size in range(1, input_size + 1):
17             for stride in range(1, input_size + 1):
18                 if cond(depth, kernel_size, stride):
19                     n = (input_size - kernel_size) // stride + 1
20                     if (n - 1) * stride + kernel_size == input_size:
21                         q = conv_chain(n, output_size, depth - 1, cond)
22                         r += [ [ (kernel_size, stride) ] + u for u in q ]
23         return r
24
25 ######################################################################
26
27 if __name__ == "__main__":
28
29     # Example
30
31     c = conv_chain(
32         input_size = 64, output_size = 8,
33         depth = 5,
34         cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d < 3)
35     )
36
37     x = torch.rand(1, 1, 64)
38
39     for m in c:
40         model = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ])
41         print(model)
42         print(x.size(), model(x).size())
43
44 ######################################################################