Update
[pytorch] / 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 ######################################################################