######################################################################
+
def conv_chain(input_size, output_size, remain_depth, cond):
if remain_depth == 0:
if input_size == output_size:
- return [ [ ] ]
+ return [[]]
else:
- return [ ]
+ return []
else:
- r = [ ]
+ r = []
for kernel_size in range(1, input_size + 1):
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:
+ 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 ]
+ r += [[(kernel_size, stride)] + u for u in q]
return r
+
######################################################################
if __name__ == "__main__":
-
import torch
from torch import nn
# Example
c = conv_chain(
- input_size = 64, output_size = 8,
- remain_depth = 5,
+ 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)
+ cond=lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2),
)
x = torch.rand(1, 1, 64)
for m in c:
- model = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ])
+ model = nn.Sequential(*[nn.Conv1d(1, 1, l[0], l[1]) for l in m])
print(model)
print(x.size(), model(x).size())