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