1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
| def conv_layer(chann_in, chann_out, k_size, p_size):
layer = tnn.Sequential(
tnn.Conv2d(chann_in, chann_out, kernel_size=k_size, padding=p_size),
tnn.BatchNorm2d(chann_out),
tnn.ReLU()
)
return layer
def vgg_conv_block(in_list, out_list, k_list, p_list, pooling_k, pooling_s):
layers = [ conv_layer(in_list[i], out_list[i], k_list[i], p_list[i]) for i in range(len(in_list)) ]
layers += [ tnn.MaxPool2d(kernel_size = pooling_k, stride = pooling_s)]
return tnn.Sequential(*layers)
def vgg_fc_layer(size_in, size_out):
layer = tnn.Sequential(
tnn.Linear(size_in, size_out),
tnn.BatchNorm1d(size_out),
tnn.ReLU()
)
return layer
|
