Merge pull request #3 from aidos-lab/mantra-dev-er

Generalized function signature for models.

models/GAT.py

@@ -37,7 +37,8 @@ def __init__(self, config: GATConfig):
             config.hidden_channels, config.out_channels
         )
 
-    def forward(self, x, edge_index, batch):
+    def forward(self, batch):
+        x, edge_index, batch = batch.x, batch.edge_index, batch.batch
         # 1. Obtain node embeddings
         x = self.gat_input(x, edge_index)
         for layer in self.hidden_layers:

models/GCN.py

@@ -36,7 +36,8 @@ def __init__(self, config: GCNConfig):
             config.hidden_channels, config.out_channels
         )
 
-    def forward(self, x, edge_index, batch):
+    def forward(self, batch):
+        x, edge_index, batch = batch.x, batch.edge_index, batch.batch
         # 1. Obtain node embeddings
         x = self.conv_input(x, edge_index)
         for layer in self.hidden_layers:

models/MLP.py

@@ -35,11 +35,12 @@ def __init__(
             config.num_hidden_neurons, config.out_channels
         )
 
-    def forward(self, x, edge_index, signal_belongings):
+    def forward(self, batch):
+        x, _, batch = batch.x, batch.edge_index, batch.batch
         x = self.input_layer(x)
         x = nn.functional.relu(x)
         for hidden_layer in self.hidden_layers:
             x = hidden_layer(x)
             x = nn.functional.relu(x)
         x = self.output_layer(x)
-        return pool.global_mean_pool(x, signal_belongings)
+        return pool.global_mean_pool(x, batch)

models/TAG.py

@@ -32,7 +32,8 @@ def __init__(
             config.hidden_channels, config.out_channels
         )
 
-    def forward(self, x, edge_index, batch):
+    def forward(self, batch):
+        x, edge_index, batch = batch.x, batch.edge_index, batch.batch
         # 1. Obtain node embeddings
         x = self.conv_input(x, edge_index)
         for layer in self.hidden_layers:

models/TransfConv.py

@@ -33,7 +33,8 @@ def __init__(
             config.hidden_channels, config.out_channels
         )
 
-    def forward(self, x, edge_index, batch):
+    def forward(self, batch):
+        x, edge_index, batch = batch.x, batch.edge_index, batch.batch
         # 1. Obtain node embeddings
         x = self.conv_input(x, edge_index)
         for layer in self.hidden_layers:

models/base.py

@@ -27,13 +27,22 @@ def __init__(
         self.model = model
         self.learning_rate = learning_rate
 
-    def forward(self, x, edge_index, batch):
-        x = self.model(x, edge_index, batch)
+    def forward(self, batch):
+        x = self.model(batch)
         return x
 
     def general_step(self, batch, batch_idx, step: str):
-        batch_len = len(batch.y)
-        x_hat = self(batch.x, batch.edge_index, batch.batch)
+
+        # This is rather ugly, open to better solutions,
+        # but torch_geometric and the toponetx dl have rather different
+        # signatures.
+        if hasattr(batch, "batch"):
+            batch_len = batch.batch.max() + 1
+        else:
+            batch_len = batch[-1]
+
+        # Generalizing to accomodate for the different signatures.
+        x_hat = self(batch)
         # Squeeze x_hat to match the shape of y
         x_hat = x_hat.squeeze()
         loss = self.loss_fn(x_hat, batch.y)