✨ add MLP

phishGNN/models/__init__.py

@@ -6,3 +6,4 @@
 from .graphsage import GraphSAGE
 from .cluster_gcn import ClusterGCN
 from .mem_pool import MemPool
+from .mlp import MLP

phishGNN/models/mlp.py

@@ -0,0 +1,39 @@
+import torch
+from torch import nn
+from torch_geometric.nn import global_mean_pool
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        in_channels=None,
+        hidden_channels=None,
+        out_channels=None,
+        pooling_fn=global_mean_pool,
+        device=None,
+        nb_layers=None,
+    ):
+        super(MLP, self).__init__()
+
+        self.pooling_fn = pooling_fn
+        self.hidden_channels = hidden_channels
+        self.device = device
+        self.to(device)
+
+        self.flatten = nn.Flatten()
+        self.lin = nn.Linear(2 * hidden_channels, hidden_channels)
+        self.lin_out = nn.Linear(hidden_channels, out_channels)
+        self.relu = nn.ReLU()
+
+
+    def forward(self, x, edge_index, batch):
+        x = x.float()
+        x = self.flatten(x)
+        x = nn.Linear(x.shape[1], 2 * self.hidden_channels)(x)
+        x = self.relu(x)
+        x = self.lin(x)
+        x = self.relu(x)
+
+        x = self.pooling_fn(x, batch)
+        x = self.lin_out(x)
+        return x

phishGNN/training.py

@@ -11,7 +11,7 @@
 from torch_geometric.loader import DataLoader
 
 from visualization import visualize, plot_embeddings
-from models import GCN_2, GCN_3, GIN, GAT, GraphSAGE, ClusterGCN, MemPool
+from models import GCN_2, GCN_3, GIN, GAT, MLP, GraphSAGE, ClusterGCN, MemPool
 from utils.utils import mean_std_error
 
 
@@ -66,6 +66,7 @@ def train(should_plot_embeddings: bool):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     models = [
+        MLP,
         GAT,
         GIN,
         GCN_2,