File size: 1,596 Bytes
0da959e |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GNN_MD(torch.nn.Module):
def __init__(self, num_features, hidden_dim):
super(GNN_MD, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.bn1 = nn.BatchNorm1d(hidden_dim)
self.conv2 = GCNConv(hidden_dim, hidden_dim*2)
self.bn2 = nn.BatchNorm1d(hidden_dim*2)
self.conv3 = GCNConv(hidden_dim*2, hidden_dim*4)
self.bn3 = nn.BatchNorm1d(hidden_dim*4)
self.conv4 = GCNConv(hidden_dim*4, hidden_dim*4)
self.bn4 = nn.BatchNorm1d(hidden_dim*4)
self.conv5 = GCNConv(hidden_dim*4, hidden_dim*8)
self.bn5 = nn.BatchNorm1d(hidden_dim*8)
self.fc1 = nn.Linear(hidden_dim*8, hidden_dim*4)
self.fc2 = nn.Linear(hidden_dim*4, 1)
def forward(self, data):
x = self.conv1(data.x, data.edge_index, data.edge_attr.view(-1))
x = F.relu(x)
x = self.bn1(x)
x = self.conv2(x, data.edge_index, data.edge_attr.view(-1))
x = F.relu(x)
x = self.bn2(x)
x = self.conv3(x, data.edge_index, data.edge_attr.view(-1))
x = F.relu(x)
x = self.bn3(x)
x = self.conv4(x, data.edge_index, data.edge_attr.view(-1))
x = self.bn4(x)
x = F.relu(x)
x = self.conv5(x, data.edge_index, data.edge_attr.view(-1))
x = self.bn5(x)
#x = global_add_pool(x, x.batch)
x = F.relu(x)
x = F.relu(self.fc1(x))
x = F.dropout(x, p=0.25)
return self.fc2(x).view(-1) |