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
  11 def conv_chain(input_size, output_size, remain_depth, cond):
  12     if remain_depth == 0:
  13         if input_size == output_size:
  14             return [[]]
  15         else:
  16             return []
  17     else:
  18         r = []
  19         for kernel_size in range(1, input_size + 1):
  20             for stride in range(1, input_size):
  21                 if cond(remain_depth, kernel_size, stride):
  22                     n = (input_size - kernel_size) // stride + 1
  23                     if (
  24                         n >= output_size
  25                         and (n - 1) * stride + kernel_size == input_size
  26                     ):
  27                         q = conv_chain(n, output_size, remain_depth - 1, cond)
  28                         r += [[(kernel_size, stride)] + u for u in q]
  29         return r
  30
  31
  32 ######################################################################
  33
  34 if __name__ == "__main__":
  35     import torch
  36     from torch import nn
  37
  38     # Example
  39
  40     c = conv_chain(
  41         input_size=64,
  42         output_size=8,
  43         remain_depth=5,
  44         # We want kernels smaller than 4, strides smaller than the
  45         # kernels, and strides of 1 except in the two last layers
  46         cond=lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2),
  47     )
  48
  49     x = torch.rand(1, 1, 64)
  50
  51     for m in c:
  52         model = nn.Sequential(*[nn.Conv1d(1, 1, l[0], l[1]) for l in m])
  53         print(model)
  54         print(x.size(), model(x).size())
  55
  56 ######################################################################