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(kernel_size, stride):
  19                     n = (input_size - kernel_size) // stride
  20                     if n * stride + kernel_size == input_size:
  21                         q = conv_chain(n + 1, 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 k, s: k <= 4 and s <= 2 and s <= k//2
  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 ######################################################################