reduce search space

kumo-ai · Nov 4, 2024 · 8faa059 · 8faa059
1 parent 8ef8111
commit 8faa059
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/benchmark/tgt_ijcai_benchmark.py b/benchmark/tgt_ijcai_benchmark.py
@@ -230,18 +230,18 @@ def create_edge(data, behavior, beh_idx, pkey_name, pkey_idx):
     model_cls = IDGNN
 elif args.model in ["contextgnn", "shallowrhsgnn"]:
     model_search_space = {
-        "encoder_channels": [32, 64, 128, 256, 512],
+        "encoder_channels": [32, 64, 128, 256],
         "encoder_layers": [2, 4, 8],
-        "channels": [32, 64, 128, 256, 512],
-        "embedding_dim": [32, 64, 128, 256, 512],
+        "channels": [32, 64, 128, 256],
+        "embedding_dim": [32, 64, 128, 256],
         "norm": ["layer_norm", "batch_norm"],
         "rhs_emb_mode": [
             RHSEmbeddingMode.FUSION, RHSEmbeddingMode.FEATURE,
             RHSEmbeddingMode.LOOKUP
         ]
     }
     train_search_space = {
-        "batch_size": [256, 512, 1024],
+        "batch_size": [128, 256, 512],
         "base_lr": [0.001, 0.01],
         "gamma_rate": [0.8, 1.],
     }
@@ -542,10 +542,8 @@ def main_gnn() -> None:
     # Save results
     if args.result_path != "":
         os.makedirs(args.result_path, exist_ok=True)
-        torch.save(
-            result_dict,
-            osp.join(args.result_path,
-                     f"{args.dataset}_{args.task}_{args.model}"))
+        torch.save(result_dict,
+                   osp.join(args.result_path, f"ijcai_{args.model}"))
 
 
 if __name__ == "__main__":