Decent model

src/models.py

@@ -98,6 +98,8 @@ def __init__(self, partition_configs):
         self.final_lin_transform = nn.Linear(self.n_partitions, 1, bias=False)
         with torch.no_grad():
             self.final_lin_transform.weight.copy_(torch.ones(self.n_partitions))
+        for p in self.final_lin_transform.parameters():
+            p.requires_grad = False
 
     def init_conv_layers(self):
         """Initialize convolutional layers."""
@@ -110,7 +112,7 @@ def init_conv_layers(self):
                     gnn.CGConv(
                         channels=self.conv_size[i],
                         dim=self.num_edge_features[i],
-                        batch_norm=False,
+                        batch_norm=True,
                     ),
                     nn.LeakyReLU(inplace=True),
                 ]

src/train.py

@@ -367,6 +367,7 @@ def predict(self, dataset, indices, return_targets=False):
             predictions[i] = self.standardizer.restore(pred_dict["output"].cpu().detach())
             conts_std = pred_dict["contributions"].cpu().detach()
             contributions[i, :] = self.standardizer.restore_cont(conts_std).numpy().flatten()
+            #contributions[i, :] = pred_dict["contributions"].cpu().detach().numpy().flatten()
 
         predictions_dict = {
             "targets": targets,

workflows/basic_train_val_test.py

@@ -45,14 +45,14 @@
     "metric_function": "mae",
     "learning_rate": 0.001,
     "optimizer": "adam",
-    "lr_milestones": [100]
+    "lr_milestones": [75]
 }
 partition_configs = [
         {
-            "n_conv": 5,
-            "n_hidden": 2,
-            "hidden_size": 50,
-            "conv_size": 50,
+            "n_conv": 3,
+            "n_hidden": 1,
+            "hidden_size": 20,
+            "conv_size": 20,
             "dropout": 0.1,
             "num_node_features": dataset[0][0].num_node_features,
             "num_edge_features": dataset[0][0].num_edge_features,
@@ -82,13 +82,13 @@
 
 model_path = REPO_PATH / "trained_models" / "S_binary_calcs" 
 model = Model(global_config, partition_configs, model_path)
-#model.init_standardizer([dataset[i][0].y for i in sample_idx["train"]])
-#results_dict = model.train(200, dataloader_dict, verbose=True)
-#print(f"Test metric: {results_dict['metric']['test']}")
+model.init_standardizer([dataset[i][0].y for i in sample_idx["train"]])
+results_dict = model.train(100, dataloader_dict, verbose=True)
+print(f"Test metric: {results_dict['metric']['test']}")
 
 # Load model
 model.load(best_status=True)
 
 # Make predictions on a structure
-pred_dict = model.predict(dataset, [200], return_targets=True)
+pred_dict = model.predict(dataset, [0, 100, 200, 500], return_targets=True)
 print(pred_dict)