ae_size.py

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