Minor changes in dataset.py and run_exp.py

david-baek · david-baek · commit 5bdaaf4c24eb · 2024-12-27T23:46:18.000-05:00
diff --git a/notebooks/plot_runs.ipynb b/notebooks/plot_runs.ipynb
diff --git a/scripts/circle_run.sh b/scripts/circle_run.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/run_exp.py --data_id circle --model_id standard_MLP
+python ../src/run_exp.py --data_id circle --model_id H_MLP
+python ../src/run_exp.py --data_id circle --model_id standard_transformer
+python ../src/run_exp.py --data_id circle --model_id H_transformer
+
diff --git a/scripts/equiv_run.sh b/scripts/equiv_run.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/run_exp.py --data_id equivalence --model_id standard_MLP
+python ../src/run_exp.py --data_id equivalence --model_id H_MLP
+python ../src/run_exp.py --data_id equivalence --model_id standard_transformer
+python ../src/run_exp.py --data_id equivalence --model_id H_transformer
+
diff --git a/scripts/family_tree_run.sh b/scripts/family_tree_run.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/run_exp.py --data_id family_tree --model_id standard_MLP
+python ../src/run_exp.py --data_id family_tree --model_id H_MLP
+python ../src/run_exp.py --data_id family_tree --model_id standard_transformer
+python ../src/run_exp.py --data_id family_tree --model_id H_transformer
+
diff --git a/scripts/greater_run.sh b/scripts/greater_run.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/run_exp.py --data_id greater --model_id standard_MLP
+python ../src/run_exp.py --data_id greater --model_id H_MLP
+python ../src/run_exp.py --data_id greater --model_id standard_transformer
+python ../src/run_exp.py --data_id greater --model_id H_transformer
+
diff --git a/scripts/lattice_run.sh b/scripts/lattice_run.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/run_exp.py --data_id lattice --model_id standard_MLP
+python ../src/run_exp.py --data_id lattice --model_id H_MLP
+python ../src/run_exp.py --data_id lattice --model_id standard_transformer
+python ../src/run_exp.py --data_id lattice --model_id H_transformer
+
diff --git a/scripts/u_circle.sh b/scripts/u_circle.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/unit_exp.py --data_id circle --model_id standard_transformer
+python ../src/unit_exp.py --data_id circle --model_id H_transformer
+python ../src/unit_exp.py --data_id circle --model_id standard_MLP
+python ../src/unit_exp.py --data_id circle --model_id H_MLP
+
diff --git a/scripts/u_equiv.sh b/scripts/u_equiv.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/unit_exp.py --data_id equivalence --model_id standard_transformer
+python ../src/unit_exp.py --data_id equivalence --model_id H_transformer
+python ../src/unit_exp.py --data_id equivalence --model_id standard_MLP
+python ../src/unit_exp.py --data_id equivalence --model_id H_MLP
+
diff --git a/scripts/u_family.sh b/scripts/u_family.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/unit_exp.py --data_id family_tree --model_id standard_transformer
+python ../src/unit_exp.py --data_id family_tree --model_id H_transformer
+python ../src/unit_exp.py --data_id family_tree --model_id standard_MLP
+python ../src/unit_exp.py --data_id family_tree --model_id H_MLP
+
diff --git a/scripts/u_greater.sh b/scripts/u_greater.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/unit_exp.py --data_id greater --model_id standard_transformer
+python ../src/unit_exp.py --data_id greater --model_id H_transformer
+python ../src/unit_exp.py --data_id greater --model_id standard_MLP
+python ../src/unit_exp.py --data_id greater --model_id H_MLP
+
diff --git a/scripts/u_lattice.sh b/scripts/u_lattice.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:1
+#SBATCH -n 16
+
+python ../src/unit_exp.py --data_id lattice --model_id standard_transformer
+python ../src/unit_exp.py --data_id lattice --model_id H_transformer
+python ../src/unit_exp.py --data_id lattice --model_id standard_MLP
+python ../src/unit_exp.py --data_id lattice --model_id H_MLP
+
diff --git a/src/run_exp.py b/src/run_exp.py
@@ -61,7 +61,7 @@
 
 ## Exp2: Metric vs Overall Dataset Size (fixed train-test split)
 print(f"Experiment 2: Metric vs Overall Dataset Size (fixed train-test split)")
-data_size_list = [100, 200, 500, 1000, 2000, 5000, 10000]
+data_size_list = [400, 500, 600, 700, 800, 900, 1000, 1100, 1200]
 for i in tqdm(range(len(data_size_list))):
     data_size = data_size_list[i]
     param_dict = {
@@ -90,14 +90,14 @@
     elif data_id == "family_tree":
         aux_info["dict_level"] = dataset['dict_level']
     elif data_id == "equivalence":
-        aux_info["mod"] = 5
+        aux_info["mod"] = 10
     elif data_id == "circle":
         aux_info["p"] = 59
     else:
         raise ValueError(f"Unknown data_id: {data_id}")
     
     if hasattr(model.embedding, 'weight'):
-        metric_dict = crystal_metric(model.embedding.weight.cpu(), data_id, aux_info)
+        metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)
     else:
         metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
 
@@ -136,14 +136,14 @@
     elif data_id == "family_tree":
         aux_info["dict_level"] = dataset['dict_level']
     elif data_id == "equivalence":
-        aux_info["mod"] = 5
+        aux_info["mod"] = 10
     elif data_id == "circle":
         aux_info["p"] = 59
     else:
         raise ValueError(f"Unknown data_id: {data_id}")
     
     if hasattr(model.embedding, 'weight'):
-        metric_dict = crystal_metric(model.embedding.weight.cpu(), data_id, aux_info)
+        metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)
     else:
         metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
 
@@ -155,6 +155,7 @@
 ## Exp4: Grokking plot: Run with different seeds
 print(f"Experiment 4: Train with different seeds")
 seed_list = np.linspace(0, 1000, 20, dtype=int)
+
 for i in tqdm(range(len(seed_list))):
     seed = seed_list[i]
     data_size = 1000
@@ -186,14 +187,14 @@
     elif data_id == "family_tree":
         aux_info["dict_level"] = dataset['dict_level']
     elif data_id == "equivalence":
-        aux_info["mod"] = 5
+        aux_info["mod"] = 10
     elif data_id == "circle":
-        aux_info["p"] = 59
+        aux_info["p"] = 31
     else:
         raise ValueError(f"Unknown data_id: {data_id}")
 
     if hasattr(model.embedding, 'weight'):
-        metric_dict = crystal_metric(model.embedding.weight.cpu(), data_id, aux_info)
+        metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)
     else:
         metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
 
diff --git a/src/unit_exp.py b/src/unit_exp.py
@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+import random
+
+from tqdm import tqdm
+
+import sys
+sys.path.append("..")
+
+import argparse
+from src.utils.driver import train_single_model
+from src.utils.visualization import visualize_embedding
+from src.utils.crystal_metric import crystal_metric
+import json
+
+data_id_choices = ["lattice", "greater", "family_tree", "equivalence", "circle"]
+model_id_choices = ["H_MLP", "standard_MLP", "H_transformer", "standard_transformer"]
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Experiment')
+    parser.add_argument('--seed', type=int, default=51, help='random seed')
+    parser.add_argument('--data_id', type=str, required=True, choices=data_id_choices, help='Data ID')
+    parser.add_argument('--model_id', type=str, required=True, choices=model_id_choices, help='Model ID')
+
+
+args = parser.parse_args()
+seed = args.seed
+data_id = args.data_id
+model_id = args.model_id
+
+data_size = 1000
+train_ratio = 0.8
+
+param_dict = {
+    'seed': seed,
+    'data_id': data_id,
+    'data_size': data_size,
+    'train_ratio': train_ratio,
+    'model_id': model_id,
+    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
+    'embd_dim': 16,
+}
+
+
+# Train the model
+print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}")
+ret_dic = train_single_model(param_dict)
+
+## Exp1: Visualize Embeddings
+print(f"Experiment 1: Visualize Embeddings")
+model = ret_dic['model']
+dataset = ret_dic['dataset']
+torch.save(model.state_dict(), f"../results/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.pt")
+
+if hasattr(model.embedding, 'weight'):
+    visualize_embedding(model.embedding.weight.cpu(), title=f"{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}", save_path=f"../results/unit_tests/emb_{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.png", dict_level = dataset['dict_level'] if 'dict_level' in dataset else None)
+else:
+    visualize_embedding(model.embedding.data.cpu(), title=f"{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}", save_path=f"../results/unit_tests/emb_{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.png", dict_level = dataset['dict_level'] if 'dict_level' in dataset else None)
+
+with open(f"../results/unit_tests/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
+        json.dump(ret_dic["results"], f, indent=4)
+
+aux_info = {}
+if data_id == "lattice":
+    aux_info["lattice_size"] = 5
+elif data_id == "greater":
+    aux_info["p"] = 200
+elif data_id == "family_tree":
+    aux_info["dict_level"] = dataset['dict_level']
+elif data_id == "equivalence":
+    aux_info["mod"] = 5
+elif data_id == "circle":
+    aux_info["p"] = 31
+else:
+    raise ValueError(f"Unknown data_id: {data_id}")
+
+if hasattr(model.embedding, 'weight'):
+    metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)
+else:
+    metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
+
+with open(f"../results/unit_tests/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.json", "w") as f:
+    json.dump(metric_dict, f, indent=4)
+
diff --git a/src/utils/FamilyTreeGenerator.py b/src/utils/FamilyTreeGenerator.py
@@ -70,7 +70,7 @@ def GenerateFamilyTree(nodes_MAX = 200, max_child_per_gen = 6, seed = 33, real_w
             cur_person = bfs_lst.pop()
             rng_num = np.random.randint(2)
 
-            if rng_num > 0: # marry with 70% probability
+            if rng_num >= 0: # marry with 70% probability
                 dict_couple_idx[cur_person] = couple_idx
                 dict_couple_idx[cur_idx] = couple_idx
                 couple_idx += 1
@@ -82,7 +82,7 @@ def GenerateFamilyTree(nodes_MAX = 200, max_child_per_gen = 6, seed = 33, real_w
                 cur_idx += 1
                 rng_num = np.random.randint(2)
 
-                if rng_num > 0: # have children with 70% probability
+                if rng_num >= 0: # have children with 70% probability
                     num_child = np.random.randint(1,max_child_per_gen+1)
                     list_sons = []
                     list_daughters = []
diff --git a/src/utils/crystal_metric.py b/src/utils/crystal_metric.py
@@ -72,7 +72,7 @@ def side_length_deviation(a, b, c, d):
     variances = pca.explained_variance_ratio_
 
     metric_dict = {
-        'metric': np.mean(deviation_arr),
+        'metric': float(np.mean(deviation_arr)),
         'variances': variances.tolist(),
     }
 
@@ -92,12 +92,17 @@ def greater_metric(reps, aux_info):
     variances = pca.explained_variance_ratio_
 
     metric_dict = {
-        'metric': np.std(diff_arr) / np.mean(diff_arr),
+        'metric': float(np.std(diff_arr) / np.mean(diff_arr)),
         'variances': variances.tolist(),
     }
     return metric_dict
 
 def family_tree_metric(reps, aux_info):
+
+    pca = PCA(n_components=min(reps.shape[0], reps.shape[1]))
+    reps = pca.fit_transform(reps)
+    reps = reps[:, :2]
+
     dict_level = aux_info['dict_level']
 
     # Group individuals by generation
@@ -129,12 +134,10 @@ def family_tree_metric(reps, aux_info):
             collinearity_by_generation[generation] = collinearity.mean()
             
 
-    pca = PCA(n_components=min(reps.shape[0], reps.shape[1]))
-    emb_pca = pca.fit_transform(reps)
     variances = pca.explained_variance_ratio_
 
     metric_dict = {
-        'metric': 1 - np.mean([collinearity for collinearity in collinearity_by_generation.values() if not np.isnan(collinearity)]),
+        'metric': float(1 - np.mean([collinearity for collinearity in collinearity_by_generation.values() if not np.isnan(collinearity)])),
         'variances': variances.tolist(),
     }
     return metric_dict
@@ -159,7 +162,7 @@ def equivalence_metric(reps, aux_info):
 
     print(np.mean(diff_arr) , np.mean(cross_diff_arr))
     metric_dict = {
-        'metric': np.mean(diff_arr) / np.mean(cross_diff_arr),
+        'metric': float(np.mean(diff_arr) / np.mean(cross_diff_arr)),
         'variances': variances.tolist(),
     }
     return metric_dict
@@ -186,7 +189,7 @@ def circle_metric(reps, aux_info):
 
 
     metric_dict = {
-        'metric': circularity_score,
+        'metric': float(circularity_score),
         'variances': variances.tolist(),
     }
     return metric_dict
diff --git a/src/utils/dataset.py b/src/utils/dataset.py
diff --git a/src/utils/driver.py b/src/utils/driver.py
diff --git a/src/utils/model.py b/src/utils/model.py
diff --git a/src/utils/visualization.py b/src/utils/visualization.py
