Spaces:
Running
Running
File size: 1,776 Bytes
f484830 |
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 |
import torch
import torch.nn as nn
from torch_geometric.nn import GINConv, global_add_pool, global_mean_pool
import torch.nn.functional as F
import numpy as np
class GIN(torch.nn.Module):
def __init__(self, num_features, num_classes, dropout, hidden_dim=64, num_layers=5, add_or_mean="add"):
super().__init__()
self.num_layers = num_layers
self.hidden_dim = hidden_dim
self.add_or_mean = add_or_mean
self.dropout = dropout
self.conv_layers = nn.ModuleList()
# input features → hidden_dim
mlp = nn.Sequential(
nn.Linear(num_features, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim)
)
self.conv_layers.append(GINConv(mlp, train_eps=True))
# hidden GIN layers
for _ in range(num_layers - 1):
mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim)
)
self.conv_layers.append(GINConv(mlp, train_eps=True))
# Final classifier (after pooling)
self.fc = nn.Linear(hidden_dim, num_classes)
def forward(self, x, edge_index, batch):
for conv in self.conv_layers:
x = conv(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
# Pool to get graph-level representation
if self.add_or_mean == "mean":
x = global_mean_pool(x, batch)
elif self.add_or_mean == "add":
x = global_add_pool(x, batch)
x = F.dropout(x, p=0.5, training=self.training)
return self.fc(x)
|