-
Notifications
You must be signed in to change notification settings - Fork 1
/
model.py
50 lines (43 loc) · 1.48 KB
/
model.py
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
import torch
import torch.nn as nn
class Net(nn.Module):
"""
The network class (discriminator network)
"""
def __init__(self, args):
super(Net, self).__init__()
self.args = args
self.sigmoid = nn.Sigmoid()
self.net_gen()
def net_gen(self):
"""
Generates the NN
:return: Nothing is returned
"""
in_size = self.args.proj_size
mid_size = self.args.mid_size
modules_tail = [nn.Linear(in_size, 2 * mid_size, bias=True),
nn.ReLU(),
nn.Linear(2 * mid_size, mid_size, bias=True),
nn.ReLU(),
nn.Linear(mid_size, mid_size//2, bias=True),
nn.ReLU(),
nn.Linear(mid_size//2, mid_size//4, bias=True),
nn.ReLU(),
nn.Linear(mid_size//4, 1, bias=True)]
self.tail = nn.Sequential(*modules_tail)
# initializing the layers
self.tail.apply(self.weights_init)
def forward(self, x):
x = self.tail(x)
return x
def weights_init(self, m):
"""
Initializes the weights of the network
:param m: the parameters of the network
:return: Nothing returned, parameters initialized in-place
"""
classname = m.__class__.__name__
if (classname.find('Conv') != -1) or (classname.find('Linear') != -1):
m.weight.data.normal_(0.0, 0.05)
m.bias.data.fill_(0)