| 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 = F.relu(x) |
| x = F.relu(self.fc1(x)) |
| x = F.dropout(x, p=0.25) |
| return self.fc2(x).view(-1) |
