ae_size.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 import math
   9 from torch import nn
  10 from torch import Tensor
  11
  12 ######################################################################
  13
  14
  15 def minimal_input_size(w, layer_specs):
  16     assert w > 0, "The input is too small"
  17     if layer_specs == []:
  18         return w
  19     else:
  20         kernel_size, stride = layer_specs[0]
  21         v = int(math.ceil((w - kernel_size) / stride)) + 1
  22         v = minimal_input_size(v, layer_specs[1:])
  23         return (v - 1) * stride + kernel_size
  24
  25
  26 ######################################################################
  27
  28 # Dummy test
  29
  30 if __name__ == "__main__":
  31     layer_specs = [(17, 5), (5, 4), (3, 2), (3, 2)]
  32
  33     layers = []
  34
  35     for kernel_size, stride in layer_specs:
  36         layers.append(nn.Conv2d(1, 1, kernel_size, stride))
  37
  38     for kernel_size, stride in reversed(layer_specs):
  39         layers.append(nn.ConvTranspose2d(1, 1, kernel_size, stride))
  40
  41     m = nn.Sequential(*layers)
  42
  43     h = minimal_input_size(240, layer_specs)
  44     w = minimal_input_size(320, layer_specs)
  45
  46     x = Tensor(1, 1, h, w).normal_()
  47
  48     print(x.size(), m(x).size())
  49
  50 ######################################################################