-
Notifications
You must be signed in to change notification settings - Fork 0
/
models.py
37 lines (30 loc) · 1.13 KB
/
models.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
import torch.nn as nn
import torch.nn.functional as F
class Fully(nn.Module):
def __init__(self, nfeat, nhid, nclass, dropout,multi_label=False):
super(Fully, self).__init__()
self.multi_label = multi_label
layers = []
if len(nhid) == 0:
layers.append(nn.Linear(nfeat, nclass))
else:
layers.append(nn.Linear(nfeat, nhid[0]))
for i in range(len(nhid) - 1):
layers.append(nn.Linear(nhid[i], nhid[i + 1]))
if nclass > 1:
layers.append(nn.Linear(nhid[-1], nclass))
self.fc = nn.ModuleList(layers)
self.dropout = dropout
self.nclass = nclass
def forward(self, x, adj=None):
end_layer = len(self.fc) - 1 if self.nclass > 1 else len(self.fc)
for i in range(end_layer):
x = F.dropout(x, self.dropout, training=self.training)
x = self.fc[i](x)
x = F.relu(x)
classifier = self.fc[-1](x)
if self.multi_label:
classifier = classifier
else:
classifier = F.log_softmax(classifier, dim=1)
return classifier