standard_nn.py
3.84 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
"""
Utility to create simple sequential networks for classification or regression
Create feed forward network with different `hidden_sizes`
Create convolution networks with different `channels` (hidden_size), (2,2) max pooling
"""
import typing
import numpy
import torch.nn as nn
from lib.utils.torch_utils import Reshape, infer_shape
# TODO : extend to cover pooling sizes, strides etc for conv nets
def get_arch(input_shape: typing.Union[numpy.array, typing.List], output_size: int,
feed_forward: bool, hidden_sizes: typing.List[int],
kernel_size: typing.Union[typing.List[int], int] = 3,
non_linearity: typing.Union[typing.List[str], str, None] = "relu",
norm: typing.Union[typing.List[str], str, None] = None,
pooling: typing.Union[typing.List[str], str, None] = None) -> nn.Module:
# general assertions
n_layers = len(hidden_sizes)
if n_layers > 0:
if isinstance(non_linearity, list):
assert len(non_linearity) == n_layers, "non linearity list is not same as hidden size"
non_linearities = non_linearity
else:
non_linearities = [non_linearity] * n_layers
if isinstance(norm, list):
assert len(norm) == n_layers, "norm list is not same as hidden size"
norms = norm
else:
norms = [norm] * n_layers
else:
norms = []
non_linearities = []
modules = []
if feed_forward:
modules.append(Reshape())
insize = int(numpy.prod(input_shape))
for nl, no, outsize in zip(non_linearities, norms, hidden_sizes):
modules.append(nn.Linear(insize, outsize))
if nl == "relu":
modules.append(nn.ReLU())
elif nl is None:
pass
else:
raise Exception(f"non-linearity {nl} not implemented")
if no == "bn":
modules.append(nn.BatchNorm1d(outsize))
elif no is None:
pass
else:
raise Exception(f"norm {no} is not implemented")
insize = outsize
modules.append(nn.Linear(insize, output_size))
return {"net" : nn.Sequential(*modules)}
# assertion specific to convolutions
assert n_layers >= 1, "Number of layers has to be more than 1 for convolution"
if isinstance(kernel_size, list):
assert len(kernel_size) == n_layers, "kernel size is not same as hidden size"
kernel_sizes = kernel_size
else:
kernel_sizes = [kernel_size] * n_layers
if isinstance(pooling, list):
assert len(pooling) == n_layers, "pooling size is not same as hidden size"
poolings = pooling
else:
poolings = [pooling] * n_layers
# convolutional layer with 3x3 convolutions
inchannel = input_shape[0]
for nl, no, outchannel, k, p in zip(non_linearities, norms, hidden_sizes, kernel_sizes,
poolings):
modules.append(nn.Conv2d(inchannel, outchannel, kernel_size=k))
if nl == "relu":
modules.append(nn.ReLU())
elif nl is None:
pass
else:
raise Exception(f"non-linearity {nl} is not implemented")
if no == "bn":
modules.append(nn.BatchNorm2d(outchannel))
elif no is None:
pass
else:
raise Exception(f"norm {no} is not implemented")
if p == "max_pool":
modules.append(nn.MaxPool2d(2))
elif p is None:
pass
else:
raise Exception(f"pooling {p} is not implemented")
inchannel = outchannel
output_shape = infer_shape(nn.Sequential(*modules).to("cpu"), input_shape)
modules.append(Reshape())
modules.append(nn.Linear(int(numpy.prod(output_shape)), output_size))
return {"net" : nn.Sequential(*modules)}