-
Notifications
You must be signed in to change notification settings - Fork 7
/
layer.py
100 lines (79 loc) · 2.97 KB
/
layer.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
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
import math
import torch
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
import torch.nn.functional as F
class GNN_Layer(Module):
"""
Layer defined for GNN-Bet
"""
def __init__(self, in_features, out_features, bias=True):
super(GNN_Layer, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = Parameter(torch.FloatTensor(in_features, out_features))
if bias:
self.bias = Parameter(torch.FloatTensor(out_features))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
stdv = 1. / math.sqrt(self.weight.size(1))
self.weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform_(-stdv, stdv)
def forward(self, input, adj):
support = torch.mm(input, self.weight)
output = torch.spmm(adj, support)
if self.bias is not None:
return output + self.bias
else:
return output
def __repr__(self):
return self.__class__.__name__ + ' (' \
+ str(self.in_features) + ' -> ' \
+ str(self.out_features) + ')'
class GNN_Layer_Init(Module):
"""
First layer of GNN_Init, for embedding lookup
"""
def __init__(self, in_features, out_features, bias=True):
super(GNN_Layer_Init, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = Parameter(torch.FloatTensor(in_features, out_features))
if bias:
self.bias = Parameter(torch.FloatTensor(out_features))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
stdv = 1. / math.sqrt(self.weight.size(1))
self.weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform_(-stdv, stdv)
def forward(self, adj):
support = self.weight
output = torch.spmm(adj, support)
if self.bias is not None:
return output + self.bias
else:
return output
def __repr__(self):
return self.__class__.__name__ + ' (' \
+ str(self.in_features) + ' -> ' \
+ str(self.out_features) + ')'
class MLP(Module):
def __init__(self, nhid,dropout):
super(MLP,self).__init__()
self.dropout = dropout
self.linear1 = torch.nn.Linear(nhid,2*nhid)
self.linear2 = torch.nn.Linear(2*nhid,2*nhid)
self.linear3 = torch.nn.Linear(2*nhid,1)
def forward(self,input_vec,dropout):
score_temp = F.relu(self.linear1(input_vec))
score_temp = F.dropout(score_temp,self.dropout)
score_temp = F.relu(self.linear2(score_temp))
score_temp = F.dropout(score_temp,self.dropout)
score_temp = self.linear3(score_temp)
return score_temp