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