-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathconv.py
201 lines (173 loc) · 8.8 KB
/
conv.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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
import torch
from torch_geometric.nn import MessagePassing
import torch.nn.functional as F
from torch_geometric.nn import global_mean_pool, global_add_pool
from ogb.graphproppred.mol_encoder import AtomEncoder,BondEncoder
from torch_geometric.utils import degree
import math
import pdb
### GIN convolution along the graph structure
class GINConv(MessagePassing):
def __init__(self, emb_dim):
super(GINConv, self).__init__(aggr = "add")
self.mlp = torch.nn.Sequential(torch.nn.Linear(emb_dim, 2*emb_dim), torch.nn.BatchNorm1d(2*emb_dim), torch.nn.ReLU(), torch.nn.Linear(2*emb_dim, emb_dim))
self.eps = torch.nn.Parameter(torch.Tensor([0]))
self.bond_encoder = BondEncoder(emb_dim = emb_dim)
def forward(self, x, edge_index, edge_attr, edge_weight=None):
edge_embedding = self.bond_encoder(edge_attr)
out = self.mlp((1 + self.eps) *x + self.propagate(edge_index, x=x, edge_attr=edge_embedding, edge_weight=edge_weight))
return out
def message(self, x_j, edge_attr, edge_weight=None):
if edge_weight is not None:
mess = F.relu((x_j + edge_attr) * edge_weight)
else:
mess = F.relu(x_j + edge_attr)
return mess
def update(self, aggr_out):
return aggr_out
### GCN convolution along the graph structure
class GCNConv(MessagePassing):
def __init__(self, emb_dim):
super(GCNConv, self).__init__(aggr='add')
self.linear = torch.nn.Linear(emb_dim, emb_dim)
self.root_emb = torch.nn.Embedding(1, emb_dim)
self.bond_encoder = BondEncoder(emb_dim = emb_dim)
def forward(self, x, edge_index, edge_attr):
x = self.linear(x)
edge_embedding = self.bond_encoder(edge_attr)
row, col = edge_index
#edge_weight = torch.ones((edge_index.size(1), ), device=edge_index.device)
deg = degree(row, x.size(0), dtype = x.dtype) + 1
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
return self.propagate(edge_index, x=x, edge_attr = edge_embedding, norm=norm) + F.relu(x + self.root_emb.weight) * 1./deg.view(-1,1)
def message(self, x_j, edge_attr, norm):
return norm.view(-1, 1) * F.relu(x_j + edge_attr)
def update(self, aggr_out):
return aggr_out
# ### GNN to generate node embedding
# class GNN_node(torch.nn.Module):
# def __init__(self, num_layer, emb_dim, drop_ratio = 0.5, JK = "last", residual = False, gnn_type = 'gin'):
# super(GNN_node, self).__init__()
# self.num_layer = num_layer
# self.drop_ratio = drop_ratio
# self.JK = JK
# ### add residual connection or not
# self.residual = residual
# if self.num_layer < 2:
# raise ValueError("Number of GNN layers must be greater than 1.")
# self.atom_encoder = AtomEncoder(emb_dim)
# ###List of GNNs
# self.convs = torch.nn.ModuleList()
# self.batch_norms = torch.nn.ModuleList()
# for layer in range(num_layer):
# if gnn_type == 'gin':
# self.convs.append(GINConv(emb_dim))
# elif gnn_type == 'gcn':
# self.convs.append(GCNConv(emb_dim))
# else:
# raise ValueError('Undefined GNN type called {}'.format(gnn_type))
# self.batch_norms.append(torch.nn.BatchNorm1d(emb_dim))
# def forward(self, batched_data, node_adv=None, edge_adv=None):
# x, edge_index, edge_attr, batch = batched_data.x, batched_data.edge_index, batched_data.edge_attr, batched_data.batch
# node_embedding = self.atom_encoder(x)
# if node_adv is not None:
# node_embedding *= node_adv
# h_list = [node_embedding]
# for layer in range(self.num_layer):
# h = self.convs[layer](h_list[layer], edge_index, edge_attr, edge_adv)
# h = self.batch_norms[layer](h)
# if layer == self.num_layer - 1:
# #remove relu for the last layer
# h = F.dropout(h, self.drop_ratio, training = self.training)
# else:
# h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
# if self.residual:
# h += h_list[layer]
# h_list.append(h)
# ### Different implementations of Jk-concat
# if self.JK == "last":
# node_representation = h_list[-1]
# elif self.JK == "sum":
# node_representation = 0
# for layer in range(self.num_layer + 1):
# node_representation += h_list[layer]
# return node_representation
class GNN_front(torch.nn.Module):
def __init__(self, num_layer, emb_dim, drop_ratio=0.5, gnn_type='gin'):
super(GNN_front, self).__init__()
self.num_layer = num_layer
self.drop_ratio = drop_ratio
self.atom_encoder = AtomEncoder(emb_dim)
self.virtualnode_embedding = torch.nn.Embedding(1, emb_dim)
torch.nn.init.constant_(self.virtualnode_embedding.weight.data, 0)
self.convs = torch.nn.ModuleList()
self.batch_norms = torch.nn.ModuleList()
self.mlp_virtualnode_list = torch.nn.ModuleList()
for layer in range(num_layer):
if gnn_type == 'gin':
self.convs.append(GINConv(emb_dim))
elif gnn_type == 'gcn':
self.convs.append(GCNConv(emb_dim))
else:
raise ValueError('Undefined GNN type called {}'.format(gnn_type))
self.batch_norms.append(torch.nn.BatchNorm1d(emb_dim))
for layer in range(num_layer - 1):
self.mlp_virtualnode_list.append(torch.nn.Sequential(torch.nn.Linear(emb_dim, 2*emb_dim),
torch.nn.BatchNorm1d(2*emb_dim),
torch.nn.ReLU(),
torch.nn.Linear(2*emb_dim, emb_dim),
torch.nn.BatchNorm1d(emb_dim),
torch.nn.ReLU()))
def forward(self, batched_data):
x, edge_index, edge_attr, batch = batched_data.x, batched_data.edge_index, batched_data.edge_attr, batched_data.batch
virtualnode_embedding = self.virtualnode_embedding(torch.zeros(batch[-1].item() + 1).to(edge_index.dtype).to(edge_index.device))
node_embedding = self.atom_encoder(x)
h_list = [node_embedding]
for layer in range(self.num_layer):
h_list[layer] = h_list[layer] + virtualnode_embedding[batch]
h = self.convs[layer](h_list[layer], edge_index, edge_attr)
h = self.batch_norms[layer](h)
if layer == self.num_layer - 1:
h = F.dropout(h, self.drop_ratio, training = self.training)
else:
h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
h = h + h_list[layer]
h_list.append(h)
if layer < self.num_layer - 1:
virtualnode_embedding_temp = global_add_pool(h_list[layer], batch) + virtualnode_embedding
virtualnode_embedding = virtualnode_embedding + F.dropout(self.mlp_virtualnode_list[layer](virtualnode_embedding_temp), self.drop_ratio, training = self.training)
node_representation = h_list[-1]
return node_representation
class GNN_backs(torch.nn.Module):
def __init__(self, num_layer, emb_dim, drop_ratio=0.5, gnn_type='gin'):
super(GNN_backs, self).__init__()
self.num_layer = num_layer
self.drop_ratio = drop_ratio
self.convs = torch.nn.ModuleList()
self.batch_norms = torch.nn.ModuleList()
for layer in range(num_layer):
if gnn_type == 'gin':
self.convs.append(GINConv(emb_dim))
elif gnn_type == 'gcn':
self.convs.append(GCNConv(emb_dim))
else:
raise ValueError('Undefined GNN type called {}'.format(gnn_type))
self.batch_norms.append(torch.nn.BatchNorm1d(emb_dim))
def forward(self, x, edge_index, edge_attr, batch, node_adv=None, edge_adv=None):
node_embedding = x
if node_adv is not None:
node_embedding = node_embedding * node_adv
h_list = [node_embedding]
for layer in range(self.num_layer):
h = self.convs[layer](h_list[layer], edge_index, edge_attr, edge_adv)
h = self.batch_norms[layer](h)
if layer == self.num_layer - 1:
h = F.dropout(h, self.drop_ratio, training = self.training)
else:
h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
h += h_list[layer]
h_list.append(h)
node_representation = h_list[-1]
return node_representation