Minor changes

david-baek · david-baek · commit f1e99d303472 · 2024-12-28T17:34:51.000-05:00
diff --git a/scripts/HM_circle.sh b/scripts/HM_circle.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id circle --model_id H_MLP
 
diff --git a/scripts/HM_equiv.sh b/scripts/HM_equiv.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id equivalence --model_id H_MLP
 
diff --git a/scripts/HM_family.sh b/scripts/HM_family.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id family_tree --model_id H_MLP
 
diff --git a/scripts/HM_lattice.sh b/scripts/HM_lattice.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id lattice --model_id H_MLP
 
diff --git a/scripts/HT_circle.sh b/scripts/HT_circle.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id circle --model_id H_transformer
 
diff --git a/scripts/HT_equiv.sh b/scripts/HT_equiv.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id equivalence --model_id H_transformer
 
diff --git a/scripts/HT_family.sh b/scripts/HT_family.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id family_tree --model_id H_transformer
 
diff --git a/scripts/HT_lattice.sh b/scripts/HT_lattice.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id lattice --model_id H_transformer
 
diff --git a/scripts/M_circle.sh b/scripts/M_circle.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id circle --model_id standard_MLP
 
diff --git a/scripts/M_equiv.sh b/scripts/M_equiv.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id equivalence --model_id standard_MLP
 
diff --git a/scripts/M_family.sh b/scripts/M_family.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id family_tree --model_id standard_MLP
 
diff --git a/scripts/M_lattice.sh b/scripts/M_lattice.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id lattice --model_id standard_MLP
 
diff --git a/scripts/T_circle.sh b/scripts/T_circle.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id circle --model_id standard_transformer
 
diff --git a/scripts/T_equiv.sh b/scripts/T_equiv.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id equivalence --model_id standard_transformer
 
diff --git a/scripts/T_family.sh b/scripts/T_family.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id family_tree --model_id standard_transformer
 
diff --git a/scripts/T_lattice.sh b/scripts/T_lattice.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
-#SBATCH -n 16
+#SBATCH --gres=gpu:a100:1
 
 python ../src/run_exp.py --data_id lattice --model_id standard_transformer
 
diff --git a/src/run_exp.py b/src/run_exp.py
@@ -51,6 +51,8 @@
     aux_info["mod"] = 5
 elif data_id == "circle":
     aux_info["p"] = 31
+elif data_id == "family_tree":
+    aux_info["dict_level"] = 2
 else:
     raise ValueError(f"Unknown data_id: {data_id}")
 
diff --git a/src/utils/crystal_metric.py b/src/utils/crystal_metric.py
@@ -100,7 +100,7 @@ def greater_metric(reps, aux_info):
 def family_tree_metric(reps, aux_info):
 
     dict_level = aux_info['dict_level']
-    reps = reps[:(max(dict_level.keys()) + 1)]
+    reps = reps[1:(max(dict_level.keys()) + 1)]
 
     pca = PCA(n_components=min(reps.shape[0], reps.shape[1]))
     reps = pca.fit_transform(reps)
@@ -112,7 +112,7 @@ def family_tree_metric(reps, aux_info):
     for node, generation in dict_level.items():
         if generation not in levels:
             levels[generation] = []
-        levels[generation].append(reps[node])
+        levels[generation].append(reps[node-1])
     
     # Compute one-dimensionality for each generation
     level_scores = {}
diff --git a/src/utils/dataset.py b/src/utils/dataset.py
@@ -278,14 +278,14 @@ def mod_classification_dataset(p, num, seed=0, device='cpu'):
     np.random.seed(seed)
     
     N_sample = num
-    x = np.random.choice(range(p), N_sample*2).reshape(N_sample, 2)
+    x = np.random.choice(range(p), N_sample).reshape(N_sample, 1)
 
     target = np.array([x[i,0]%5 for i in range(N_sample)])
     
     data_id = torch.from_numpy(x).to(device)
     labels = torch.from_numpy(target).to(device)
     
-    vocab_size = p+2
+    vocab_size = p
     
     dataset = {}
     dataset['data_id'] = data_id
diff --git a/src/utils/driver.py b/src/utils/driver.py
@@ -67,7 +67,7 @@ def train_single_model(param_dict: dict):
     elif data_id == "family_tree":
         dataset = family_tree_dataset_2(p=127, num=data_size, seed=seed, device=device)
     elif data_id == "equivalence":
-        input_token = 2
+        input_token = 1
         dataset = mod_classification_dataset(p=100, num=data_size, seed=seed, device=device)
     elif data_id == "circle":
         dataset = modular_addition_dataset(p=31, num=data_size, seed=seed, device=device)
