avoid redundant params in PET initialization

spozdn · Dec 26, 2023 · 35703d9 · 35703d9
1 parent fdac59a
commit 35703d9
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 42 deletions.
diff --git a/src/estimate_error.py b/src/estimate_error.py
@@ -76,16 +76,8 @@ def main():
     else:        
         loader = DataLoader(graphs, batch_size=args.batch_size, shuffle=False)
 
-    add_tokens = []
-    for _ in range(ARCHITECTURAL_HYPERS.N_GNN_LAYERS - 1):
-        add_tokens.append(ARCHITECTURAL_HYPERS.ADD_TOKEN_FIRST)
-    add_tokens.append(ARCHITECTURAL_HYPERS.ADD_TOKEN_SECOND)
-
-    model = PET(ARCHITECTURAL_HYPERS, ARCHITECTURAL_HYPERS.TRANSFORMER_D_MODEL, ARCHITECTURAL_HYPERS.TRANSFORMER_N_HEAD,
-                           ARCHITECTURAL_HYPERS.TRANSFORMER_DIM_FEEDFORWARD, ARCHITECTURAL_HYPERS.N_TRANS_LAYERS, 
-                           0.0, len(all_species), 
-                           ARCHITECTURAL_HYPERS.N_GNN_LAYERS, ARCHITECTURAL_HYPERS.HEAD_N_NEURONS, ARCHITECTURAL_HYPERS.TRANSFORMERS_CENTRAL_SPECIFIC, ARCHITECTURAL_HYPERS.HEADS_CENTRAL_SPECIFIC, 
-                           add_tokens, FITTING_SCHEME.GLOBAL_AUG).to(device)
+    model = PET(ARCHITECTURAL_HYPERS, 0.0, len(all_species),
+                FITTING_SCHEME.GLOBAL_AUG).to(device)
 
     model = PETMLIPWrapper(model, MLIP_SETTINGS.USE_ENERGIES, MLIP_SETTINGS.USE_FORCES)
 

diff --git a/src/pet.py b/src/pet.py
@@ -277,16 +277,26 @@ def forward(self, batch_dict):
         outputs = self.nn(pooled)[..., 0]
         return {"atomic_energies" : outputs}
 
-
 class PET(torch.nn.Module):
-    def __init__(self, hypers, transformer_d_model, transformer_n_head,
-                       transformer_dim_feedforward, transformer_n_layers, 
-                       transformer_dropout, n_atomic_species, 
-                       n_gnn_layers, head_n_neurons, 
-                       transformers_central_specific, 
-                       heads_central_specific, add_central_tokens, global_aug):
+    def __init__(self, hypers, transformer_dropout, n_atomic_species, 
+                  global_aug):
         super(PET, self).__init__()
         self.hypers = hypers
+        transformer_d_model = hypers.TRANSFORMER_D_MODEL
+        transformer_n_head = hypers.TRANSFORMER_N_HEAD
+        transformer_dim_feedforward = hypers.TRANSFORMER_DIM_FEEDFORWARD
+        transformer_n_layers = hypers.N_TRANS_LAYERS
+        n_gnn_layers = hypers.N_GNN_LAYERS
+        head_n_neurons = hypers.HEAD_N_NEURONS
+        transformers_central_specific = hypers.TRANSFORMERS_CENTRAL_SPECIFIC
+        heads_central_specific = hypers.HEADS_CENTRAL_SPECIFIC
+
+        add_central_tokens = []
+        for _ in range(hypers.N_GNN_LAYERS - 1):
+            add_central_tokens.append(hypers.ADD_TOKEN_FIRST)
+        add_central_tokens.append(hypers.ADD_TOKEN_SECOND)
+
+
         self.embedding = nn.Embedding(n_atomic_species + 1, transformer_d_model)
         self.global_aug = global_aug
         gnn_layers = []

diff --git a/src/single_struct_calculator.py b/src/single_struct_calculator.py
@@ -28,16 +28,8 @@ def __init__(self, path_to_calc_folder, checkpoint="best_val_rmse_both_model", d
         if MLIP_SETTINGS.USE_ENERGIES:
             self.self_contributions = np.load(self_contributions_path)
 
-        add_tokens = []
-        for _ in range(ARCHITECTURAL_HYPERS.N_GNN_LAYERS - 1):
-            add_tokens.append(ARCHITECTURAL_HYPERS.ADD_TOKEN_FIRST)
-        add_tokens.append(ARCHITECTURAL_HYPERS.ADD_TOKEN_SECOND)
-
-        model = PET(ARCHITECTURAL_HYPERS, ARCHITECTURAL_HYPERS.TRANSFORMER_D_MODEL, ARCHITECTURAL_HYPERS.TRANSFORMER_N_HEAD,
-                               ARCHITECTURAL_HYPERS.TRANSFORMER_DIM_FEEDFORWARD, ARCHITECTURAL_HYPERS.N_TRANS_LAYERS, 
-                               0.0, len(all_species), 
-                               ARCHITECTURAL_HYPERS.N_GNN_LAYERS, ARCHITECTURAL_HYPERS.HEAD_N_NEURONS, ARCHITECTURAL_HYPERS.TRANSFORMERS_CENTRAL_SPECIFIC, ARCHITECTURAL_HYPERS.HEADS_CENTRAL_SPECIFIC, 
-                               add_tokens, FITTING_SCHEME.GLOBAL_AUG).to(device)
+        model = PET(ARCHITECTURAL_HYPERS, 0.0, len(all_species),
+                FITTING_SCHEME.GLOBAL_AUG).to(device)
 
         model = PETMLIPWrapper(model, MLIP_SETTINGS.USE_ENERGIES, MLIP_SETTINGS.USE_FORCES)
         if FITTING_SCHEME.MULTI_GPU and torch.cuda.is_available():

diff --git a/src/train_model.py b/src/train_model.py
@@ -95,17 +95,8 @@ def seed_worker(worker_id):
         train_loader = DataLoader(train_graphs, batch_size=FITTING_SCHEME.STRUCTURAL_BATCH_SIZE, shuffle=True, worker_init_fn=seed_worker, generator=g)
         val_loader = DataLoader(val_graphs, batch_size = FITTING_SCHEME.STRUCTURAL_BATCH_SIZE, shuffle = False, worker_init_fn=seed_worker, generator=g)
 
-
-    add_tokens = []
-    for _ in range(ARCHITECTURAL_HYPERS.N_GNN_LAYERS - 1):
-        add_tokens.append(ARCHITECTURAL_HYPERS.ADD_TOKEN_FIRST)
-    add_tokens.append(ARCHITECTURAL_HYPERS.ADD_TOKEN_SECOND)
-
-    model = PET(ARCHITECTURAL_HYPERS, ARCHITECTURAL_HYPERS.TRANSFORMER_D_MODEL, ARCHITECTURAL_HYPERS.TRANSFORMER_N_HEAD,
-                           ARCHITECTURAL_HYPERS.TRANSFORMER_DIM_FEEDFORWARD, ARCHITECTURAL_HYPERS.N_TRANS_LAYERS, 
-                           0.0, len(all_species), 
-                           ARCHITECTURAL_HYPERS.N_GNN_LAYERS, ARCHITECTURAL_HYPERS.HEAD_N_NEURONS, ARCHITECTURAL_HYPERS.TRANSFORMERS_CENTRAL_SPECIFIC, ARCHITECTURAL_HYPERS.HEADS_CENTRAL_SPECIFIC, 
-                           add_tokens, FITTING_SCHEME.GLOBAL_AUG).to(device)
+    model = PET(ARCHITECTURAL_HYPERS, 0.0, len(all_species),
+                FITTING_SCHEME.GLOBAL_AUG).to(device)
 
     model = PETMLIPWrapper(model, MLIP_SETTINGS.USE_ENERGIES, MLIP_SETTINGS.USE_FORCES)
     if FITTING_SCHEME.MULTI_GPU and torch.cuda.is_available():
@@ -171,12 +162,8 @@ def func_lr_scheduler(epoch):
         multiplication_rmse_model_keeper = ModelKeeper()
         multiplication_mae_model_keeper = ModelKeeper()
 
-    best_val_mae = None
-    best_val_model = None
     pbar = tqdm(range(FITTING_SCHEME.EPOCH_NUM))
 
-
-
     for epoch in pbar:
 
         model.train(True)