-
Notifications
You must be signed in to change notification settings - Fork 0
/
gnn.py
74 lines (61 loc) · 3.13 KB
/
gnn.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
import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing
from torch_geometric.nn import global_mean_pool, GINConv
import torch.nn.functional as F
# from conv_base import GNN_node_Virtualnode
from conv_mol import vGINMolHeadEncoder, GINMolHeadEncoder
import pdb
class GINEncoder(torch.nn.Module):
def __init__(self, num_layer, in_dim, emb_dim):
super(GINEncoder, self).__init__()
self.num_layer = num_layer
self.in_dim = in_dim
self.emb_dim = emb_dim
self.dropout_rate = 0.5
self.relu1 = nn.ReLU()
self.relus = nn.ModuleList([nn.ReLU() for _ in range(num_layer - 1)])
self.batch_norm1 = nn.BatchNorm1d(emb_dim)
self.batch_norms = nn.ModuleList([nn.BatchNorm1d(emb_dim) for _ in range(num_layer - 1)])
self.dropout1 = nn.Dropout(self.dropout_rate)
self.dropouts = nn.ModuleList([nn.Dropout(self.dropout_rate) for _ in range(num_layer - 1)])
self.conv1 = GINConv(nn.Sequential(nn.Linear(in_dim, 2 * emb_dim),
nn.BatchNorm1d(2 * emb_dim), nn.ReLU(),
nn.Linear(2 * emb_dim, emb_dim)))
self.convs = nn.ModuleList([GINConv(nn.Sequential(nn.Linear(emb_dim, 2 * emb_dim),
nn.BatchNorm1d(2 * emb_dim), nn.ReLU(),
nn.Linear(2 * emb_dim, emb_dim)))for _ in range(num_layer - 1)])
def forward(self, batched_data):
x, edge_index, batch = batched_data.x, batched_data.edge_index, batched_data.batch
post_conv = self.dropout1(self.relu1(self.batch_norm1(self.conv1(x, edge_index))))
for i, (conv, batch_norm, relu, dropout) in enumerate(
zip(self.convs, self.batch_norms, self.relus, self.dropouts)):
post_conv = batch_norm(conv(post_conv, edge_index))
if i != len(self.convs) - 1:
post_conv = relu(post_conv)
post_conv = dropout(post_conv)
return post_conv
class GINNet(torch.nn.Module):
def __init__(self, num_class, dataset, num_layer, in_dim=None, emb_dim=300):
super(GINNet, self).__init__()
self.dataset = dataset
self.num_layer = num_layer
self.in_dim = in_dim
self.emb_dim = emb_dim
self.num_class = num_class
if dataset in ["motif", "cmnist"]:
self.gnn_node = GINEncoder(num_layer, in_dim, emb_dim)
else:
self.gnn_node = vGINMolHeadEncoder(num_layer, emb_dim)
self.pool = global_mean_pool
self.classifier = torch.nn.Linear(emb_dim, self.num_class)
def forward(self, batched_data, return_feature=False):
x, edge_index, edge_attr, batch = batched_data.x, batched_data.edge_index, batched_data.edge_attr, batched_data.batch
if self.dataset in ["motif", "cmnist"]:
h_node = self.gnn_node(batched_data)
else:
h_node = self.gnn_node(x, edge_index, edge_attr, batch)
h_graph = self.pool(h_node, batched_data.batch)
if return_feature:
return h_graph
return self.classifier(h_graph)