Merge pull request #86 from ParthaPratimBanik/main

fix torch.mean(auc) and updating model init in utils.py
pygod-team · Aug 17, 2023 · 2508cd1 · 2508cd1
2 parents 987776c + 5563a1c
commit 2508cd1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/benchmark/main.py b/benchmark/main.py
@@ -32,6 +32,10 @@ def main(args):
         ap.append(eval_average_precision(y, score))
         rec.append(eval_recall_at_k(y, score, k))
 
+    auc = torch.tensor(auc)
+    ap = torch.tensor(ap)
+    rec = torch.tensor(rec)
+
     print(args.dataset + " " + model.__class__.__name__ + " " +
           "AUC: {:.4f}±{:.4f} ({:.4f})\t"
           "AP: {:.4f}±{:.4f} ({:.4f})\t"

diff --git a/benchmark/utils.py b/benchmark/utils.py
@@ -34,11 +34,7 @@ def init_model(args):
     else:
         hid_dim = [32, 64, 128, 256]
 
-    if args.dataset[:3] == 'inj' or args.dataset[:3] == 'gen':
-        # auto balancing on injected dataset
-        alpha = [None]
-    else:
-        alpha = [0.8, 0.5, 0.2]
+    alpha = [0.8, 0.5, 0.2]
 
     if model_name == "adone":
         return AdONE(hid_dim=choice(hid_dim),
@@ -70,7 +66,7 @@ def init_model(args):
                      lr=choice(lr),
                      epoch=epoch,
                      gpu=gpu,
-                     alpha=choice(alpha),
+                     weight=choice(alpha),
                      batch_size=batch_size,
                      num_neigh=num_neigh)
     elif model_name == 'dominant':
@@ -80,7 +76,7 @@ def init_model(args):
                         lr=choice(lr),
                         epoch=epoch,
                         gpu=gpu,
-                        alpha=choice(alpha),
+                        weight=choice(alpha),
                         batch_size=batch_size,
                         num_neigh=num_neigh)
     elif model_name == 'done':
@@ -100,7 +96,7 @@ def init_model(args):
                     lr=choice(lr),
                     epoch=epoch,
                     gpu=gpu,
-                    alpha=choice(alpha),
+                    weight=choice(alpha),
                     batch_size=batch_size,
                     num_neigh=num_neigh)
     elif model_name == 'gcnae':