From c8accc41236893bb2f840e59a61084861065ecd6 Mon Sep 17 00:00:00 2001
From: Gaurav S Deshmukh <deshmukg@gilbreth-fe00.rcac.purdue.edu>
Date: Mon, 25 Sep 2023 16:24:58 -0400
Subject: [PATCH] Decent model

---
 src/models.py                     |  4 +++-
 src/train.py                      |  1 +
 workflows/basic_train_val_test.py | 18 +++++++++---------
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/models.py b/src/models.py
index e880e7f..97f6cbc 100644
--- a/src/models.py
+++ b/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),
                 ]
diff --git a/src/train.py b/src/train.py
index 80dcfdb..b676420 100644
--- a/src/train.py
+++ b/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,
diff --git a/workflows/basic_train_val_test.py b/workflows/basic_train_val_test.py
index 44314dc..0aba6c5 100644
--- a/workflows/basic_train_val_test.py
+++ b/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)
\ No newline at end of file