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 layer_specs = [ (11, 5), (5, 2), (3, 2), (3, 2) ]
  22
  23 layers = []
  24 for l in layer_specs:
  25     layers.append(nn.Conv2d(1, 1, l[0], l[1]))
  26
  27 for l in reversed(layer_specs):
  28     layers.append(nn.ConvTranspose2d(1, 1, l[0], l[1]))
  29
  30 m = nn.Sequential(*layers)
  31
  32 h = minimal_input_size(240, layer_specs)
  33 w = minimal_input_size(320, layer_specs)
  34
  35 x = Tensor(1, 1, h, w).normal_()
  36
  37 print(x.size(), m(x).size())