New experiment testing exponent, saving config files, option to save images during training to make a video

Author: Centrattic

notebooks/permutation_group.ipynb

@@ -15,6 +15,9 @@
 from src.utils.crystal_metric import crystal_metric
 import json
 
+import os
+from datetime import datetime
+
 data_id_choices = ["lattice", "greater", "family_tree", "equivalence", "circle", "permutation"]
 model_id_choices = ["H_MLP", "standard_MLP", "H_transformer", "standard_transformer"]
 if __name__ == '__main__':
@@ -23,15 +26,21 @@
     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')
 
-results_root = "results_embd_n"
-
 args = parser.parse_args()
 seed = args.seed
 data_id = args.data_id
 model_id = args.model_id
 
+## ------------------------ CONFIG -------------------------- ##
+
 data_size = 1000
 train_ratio = 0.8
+embd_dim = 16
+
+lr = 0.002
+weight_decay = 0.01
+
+n_exp=embd_dim
 
 param_dict = {
     'seed': seed,
@@ -40,9 +49,24 @@
     'train_ratio': train_ratio,
     'model_id': model_id,
     'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
-    'embd_dim': 16,
+    'embd_dim': embd_dim,
+    'n_exp': n_exp,
+    'lr': lr,
+    'weight_decay':weight_decay
 }
 
+results_root = "../results_test"
+
+current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+results_root = f"{results_root}/{current_datetime}"
+os.mkdir(results_root)
+
+param_dict_json = {k: v for k, v in param_dict.items() if k != 'device'} #  since torch.device is not JSON serializable
+
+
+with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}_config.json", "w") as f:
+    json.dump(param_dict_json, f, indent=4)
+
 aux_info = {}
 if data_id == "lattice":
     aux_info["lattice_size"] = 5
@@ -59,20 +83,20 @@
 else:
     raise ValueError(f"Unknown data_id: {data_id}")
 
-# # 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)
+# Train the model
+print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}, n_exp {n_exp}, embd_dim {embd_dim}")
+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_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.pt")
+## Exp1: Visualize Embeddings
+print(f"Experiment 1: Visualize Embeddings")
+model = ret_dic['model']
+dataset = ret_dic['dataset']
+torch.save(model.state_dict(), f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.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_root}/emb_{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.png", dict_level = dataset['dict_level'] if 'dict_level' in dataset else None, color_dict = False if data_id == "permutation" else True, adjust_overlapping_text = False)
-# else:
-#     visualize_embedding(model.embedding.data.cpu(), title=f"{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}", save_path=f"../{results_root}/emb_{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.png", dict_level = dataset['dict_level'] if 'dict_level' in dataset else None, color_dict = False if data_id == "permutation" else True, adjust_overlapping_text = False)
+if hasattr(model.embedding, 'weight'):
+    visualize_embedding(model.embedding.weight.cpu(), title=f"{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}", save_path=f"{results_root}/emb_{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.png", dict_level = dataset['dict_level'] if 'dict_level' in dataset else None, color_dict = False if data_id == "permutation" else True, adjust_overlapping_text = False)
+else:
+    visualize_embedding(model.embedding.data.cpu(), title=f"{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}", save_path=f"{results_root}/emb_{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.png", dict_level = dataset['dict_level'] if 'dict_level' in dataset else None, color_dict = False if data_id == "permutation" else True, adjust_overlapping_text = False)
 
 
 # ## Exp2: Metric vs Overall Dataset Size (fixed train-test split)
@@ -87,15 +111,19 @@
 #         'train_ratio': train_ratio,
 #         'model_id': model_id,
 #         'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
-#         'embd_dim': 16,
+#         'embd_dim': embd_dim,
+#         'n_exp': n_exp,
+#         'lr': lr,
+#         'weight_decay':weight_decay
 #     }
-#     print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id} with train_ratio {train_ratio} and data_size {data_size}")
+
+#     print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}, n_exp {n_exp}, embd_dim {embd_dim}")
 #     ret_dic = train_single_model(param_dict)
 #     model = ret_dic['model']
 #     dataset = ret_dic['dataset']
 
-#     torch.save(model.state_dict(), f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.pt")
-#     with open(f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
+#     torch.save(model.state_dict(), f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.pt")
+#     with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}_train_results.json", "w") as f:
 #         json.dump(ret_dic["results"], f, indent=4)
 
 #     if data_id == "family_tree":
@@ -106,7 +134,7 @@
 #     else:
 #         metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
 
-#     with open(f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.json", "w") as f:
+#     with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.json", "w") as f:
 #         json.dump(metric_dict, f, indent=4)
 
 ## Exp3: Metric vs Train Fraction (fixed dataset size)
@@ -122,15 +150,18 @@
         'train_ratio': train_ratio,
         'model_id': model_id,
         'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
-        'embd_dim': 16,
+        'embd_dim': embd_dim,
+        'n_exp': n_exp,
+        'lr': lr,
+        'weight_decay':weight_decay
     }
-    print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id} with train_ratio {train_ratio} and data_size {data_size}")
+    print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}, n_exp {n_exp}, embd_dim {embd_dim}")
     ret_dic = train_single_model(param_dict)
     model = ret_dic['model']
     dataset = ret_dic['dataset']
 
-    torch.save(model.state_dict(), f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.pt")
-    with open(f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
+    torch.save(model.state_dict(), f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.pt")
+    with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}_train_results.json", "w") as f:
         json.dump(ret_dic["results"], f, indent=4)
 
     if data_id == "family_tree":
@@ -141,7 +172,7 @@
     else:
         metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
 
-    with open(f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.json", "w") as f:
+    with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}_metric.json", "w") as f:
         json.dump(metric_dict, f, indent=4)
 
 ## Exp4: Grokking plot: Run with different seeds
@@ -160,14 +191,57 @@
         'train_ratio': train_ratio,
         'model_id': model_id,
         'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
-        'embd_dim': 16,
+        'embd_dim': embd_dim,
+        'n_exp': n_exp,
+        'lr': lr,
+        'weight_decay':weight_decay
+    }
+    print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}, n_exp {n_exp}, embd_dim {embd_dim}")
+    ret_dic = train_single_model(param_dict)
+    model = ret_dic['model']
+    dataset = ret_dic['dataset']
+    torch.save(model.state_dict(), f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.pt")
+    with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}_train_results.json", "w") as f:
+        json.dump(ret_dic["results"], f, indent=4)
+
+    if data_id == "family_tree":
+        aux_info["dict_level"] = dataset['dict_level']
+
+    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_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.json", "w") as f:
+        json.dump(metric_dict, f, indent=4)
+
+#Exp5: N Exponent value plot: Run with different n values, plot test accuracy vs. and explained variance vs.
+
+print(f"Experiment 5: Train with different exponent values")
+n_list = np.arange(1, 17, dtype=int)
+
+for i in tqdm(range(len(n_list))):
+    n_exp = n_list[i]
+    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': embd_dim,
+        'n_exp': n_exp
     }
-    print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id} with train_ratio {train_ratio} and data_size {data_size}")
+    print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}, n_exp {n_exp}, embd_dim {embd_dim}")
+    
     ret_dic = train_single_model(param_dict)
     model = ret_dic['model']
     dataset = ret_dic['dataset']
-    torch.save(model.state_dict(), f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.pt")
-    with open(f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
+    torch.save(model.state_dict(), f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.pt")
+    with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}_train_results.json", "w") as f:
         json.dump(ret_dic["results"], f, indent=4)
 
     if data_id == "family_tree":
@@ -178,6 +252,6 @@
     else:
         metric_dict = crystal_metric(model.embedding.data.cpu(), data_id, aux_info)
 
-    with open(f"../{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.json", "w") as f:
+    with open(f"{results_root}/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_{n_exp}.json", "w") as f:
         json.dump(metric_dict, f, indent=4)
-    
+

src/run_exp.py

@@ -69,7 +69,7 @@ def permutation_group_dataset(p, num, seed=0, device='cpu'):
     torch.manual_seed(seed)
     np.random.seed(seed)
 
-    perms = list(itertools.permutations(range(4)))
+    perms = list(itertools.permutations(range(p)))
     num_perms = len(perms)
 
     perm_dict = dict(enumerate(perms))
@@ -79,13 +79,9 @@ def permutation_group_dataset(p, num, seed=0, device='cpu'):
 
     data_id = [[perms[int(i)], perms[int(j)]] for i, j in torch.cartesian_prod(idx, idx)]
     keyed_data_id = np.array([[swapped_dict[data_id[i][0]], swapped_dict[data_id[i][1]]] for i in range(len(data_id))])
-    # data_id = np.fromiter([[tuple(perms[i]), tuple(perms[j])] for i, j in zip(idx1, idx2)], object)
-    # data_id = np.array([[perms_list[int(i)], perms_list[int(j)]] for i, j in torch.cartesian_prod(idx, idx)])
 
     labels = [tuple(np.array(perms[int(i)])[np.array(perms[int(j)])]) for i, j in torch.cartesian_prod(idx, idx)]
     keyed_labels = np.array([swapped_dict[labels[i]] for i in range(len(labels))])
-    # labels = [sum(a != b for a, b in zip(lbl, idx)) for lbl in labels]
-    # labels = np.array([sum(math.pow(10, i) * num for i, num in enumerate(reversed(tup))) for tup in labels]).astype(int)
     labels = torch.tensor(labels, dtype=torch.long, device=device)
 
     perm_vals = ["".join(np.array(perm_dict[i]).astype(str)) for i in range(len(perm_dict))]

src/utils/dataset.py

@@ -47,6 +47,9 @@ def train_single_model(param_dict: dict):
         raise ValueError("device must be provided in param_dict")
     if "embd_dim" not in param_dict:
         raise ValueError("embd_dim must be provided in param_dict")
+    if "n_exp" not in param_dict: 
+        raise ValueError("n_exp must be provided in param_dict")
+
 
     seed = param_dict['seed']
     data_id = param_dict['data_id']
@@ -55,13 +58,17 @@ def train_single_model(param_dict: dict):
     model_id = param_dict['model_id']
     device = param_dict['device']
     embd_dim = param_dict['embd_dim']
+    n_exp = param_dict['n_exp']
+
+    video = False if 'video' not in param_dict else param_dict['video']
+    lr = 0.002 if 'lr' not in param_dict else param_dict['lr']
+    weight_decay = 0.01 if 'weight_decay' not in param_dict else param_dict['weight_decay']
 
     set_seed(seed)
 
-    
     # define dataset
     input_token = 2
-    num_epochs = None
+    num_epochs = 7000
     if data_id == "lattice":
         dataset = parallelogram_dataset(p=5, dim=2, num=data_size, seed=seed, device=device)
         input_token = 3
@@ -77,28 +84,27 @@ def train_single_model(param_dict: dict):
     elif data_id=="permutation":
         dataset = permutation_group_dataset(p=4, num=data_size, seed=seed, device=device)
         if model_id == "H_transformer" or model_id == "standard_transformer":
-            num_epochs = 12750 # extra epochs to train fully
+            num_epochs = 10000 # extra epochs to train fully
     else:
         raise ValueError(f"Unknown data_id: {data_id}")
 
     dataset = split_dataset(dataset, train_ratio=train_ratio, seed=seed)
     vocab_size = dataset['vocab_size']
 
-
     # define model
     if model_id == "H_MLP":
         weight_tied = True
         hidden_size = 100
         shp = [input_token * embd_dim, hidden_size, embd_dim, vocab_size]
-        model = MLP_HS(shp=shp, vocab_size=vocab_size, embd_dim=embd_dim, input_token=input_token, weight_tied=weight_tied, seed=seed, n=np.embd_dim, init_scale=1).to(device)
+        model = MLP_HS(shp=shp, vocab_size=vocab_size, embd_dim=embd_dim, input_token=input_token, weight_tied=weight_tied, seed=seed, n=n_exp, init_scale=1).to(device)
     elif model_id == "standard_MLP":
         unembd = True
         weight_tied = True
         hidden_size = 100
         shp = [input_token * embd_dim, hidden_size, embd_dim, vocab_size]
         model = MLP(shp=shp, vocab_size=vocab_size, embd_dim=embd_dim, input_token=input_token, unembd=unembd, weight_tied=weight_tied, seed=seed, init_scale=1).to(device)
     elif model_id == "H_transformer":
-        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=2, num_layers=2, n_dist=embd_dim,seq_len=input_token, seed=seed, use_dist_layer=True, init_scale=1).to(device)
+        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=2, num_layers=2, n_dist=n_exp,seq_len=input_token, seed=seed, use_dist_layer=True, init_scale=1).to(device)
     elif model_id == "standard_transformer":
         model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=2, num_layers=2, seq_len=input_token, seed=seed, use_dist_layer=False, init_scale=1).to(device)
     else:
@@ -112,7 +118,7 @@ def train_single_model(param_dict: dict):
     test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
     ret_dic = {}
-    ret_dic["results"] = model.train(param_dict={'num_epochs': num_epochs if num_epochs else 7000, 'learning_rate': 0.002, 'train_dataloader': train_dataloader, 'test_dataloader': test_dataloader, 'device': device})
+    ret_dic["results"] = model.train(param_dict={'num_epochs': num_epochs, 'learning_rate': lr, 'weight_decay':weight_decay, 'train_dataloader': train_dataloader, 'test_dataloader': test_dataloader, 'device': device, 'video': video})
     ret_dic["model"] = model
     ret_dic["dataset"] = dataset
 

src/utils/driver.py

@@ -5,6 +5,7 @@
 import numpy as np
 import math
 from src.utils.dataset import *
+from src.utils.visualization import *
 
 import sys
 # import keyboard
@@ -23,7 +24,8 @@ def train(self, param_dict: dict):
         train_dataloader = param_dict['train_dataloader']
         test_dataloader = param_dict['test_dataloader']
         device = param_dict['device']
-        weight_decay = 0.01 if 'weight_decay' not in param_dict else param_dict['weight_decay']
+        weight_decay = param_dict['weight_decay']
+        video = False if 'video' not in param_dict else param_dict['video']
 
         verbose = True
         if 'verbose' in param_dict:
@@ -45,6 +47,12 @@ def train(self, param_dict: dict):
             # if keyboard.is_pressed('ctrl+d'):
             #     print("Manual early stopping occurring.")
             #     break
+            if video and epoch%10 == 0: # save every 10 epochs
+                if hasattr(self.embedding, 'weight'):
+                    embd = self.embedding.weight
+                else:
+                    embd = self.embedding.data
+                visualize_embedding(embd, title=f"Epoch {epoch}", save_path=f"../video_imgs/{epoch}.png", dict_level = None, color_dict = True, adjust_overlapping_text = False)
 
             train_loss = 0
             train_correct = 0
@@ -305,7 +313,7 @@ def forward(self, x):
         return logits
 
 
-def load_model_from_file(model_id, data_id, data_size = 1000, train_ratio=0.8,seed=66, embd_dim=16, device='cpu'):
+def load_model_from_file(model_id, data_id, results_root = "results",data_size = 1000, train_ratio=0.8,seed=66, embd_dim=16, device='cpu'):
 
     input_token=2
 
@@ -348,6 +356,6 @@ def load_model_from_file(model_id, data_id, data_size = 1000, train_ratio=0.8,se
     else:
         raise ValueError(f"Unknown model_id: {model_id}")
 
-    model.load_state_dict(torch.load(f"../results/{seed}_permutation_{model_id}_{data_size}_{train_ratio}.pt"))
+    model.load_state_dict(torch.load(f"../{results_root}/{seed}_permutation_{model_id}_{data_size}_{train_ratio}.pt"))
 
     return model

src/utils/model.py

@@ -43,6 +43,8 @@ def visualize_embedding(emb, title="", save_path=None, dict_level = None, color_
         adjust_text(texts, x=x, y=y, autoalign='xy', force_points=0.5, only_move = {'text':'xy'})
     if save_path:
         plt.savefig(save_path)
+    plt.show()
+    plt.close()
 
 
 def visualize_embedding_3d(emb, title="", save_path=None, dict_level = None, color_dict=True):