Updated visualization notebook

Author: david-baek

notebooks/plot_runs.ipynb

@@ -0,0 +1,10 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH -p tegmark
+#SBATCH --gres=gpu:a100:1
+
+for ARG in $(python -c "import numpy as np; print(' '.join(map(str, np.linspace(0, 1000, 20, dtype=int))))"); do
+    echo "Running with seed $ARG:"
+    python ../src/unit_exp.py --data_id circle --model_id H_transformer --seed $ARG
+    echo
+done

scripts/u_circle_new.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+#SBATCH -t 16:00:00
+#SBATCH --gres=gpu:a100:1
+
+for ARG in $(python -c "import numpy as np; print(' '.join(map(str, np.linspace(0, 1000, 20, dtype=int))))"); do
+    echo "Running with seed $ARG:"
+    python ../src/unit_exp.py --data_id family_tree --model_id H_transformer --seed $ARG
+    echo
+done

scripts/u_family_new.sh

@@ -51,15 +51,15 @@
 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")
+torch.save(model.state_dict(), f"../results/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_d=sqrtembed_1.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)
+    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}_new.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)
+    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}_new.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)
+with open(f"../results/unit_tests/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results_new.json", "w") as f:
+    json.dump(ret_dic["results"], f, indent=4)
 
 aux_info = {}
 if data_id == "lattice":
@@ -80,6 +80,6 @@
 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:
+with open(f"../results/unit_tests/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_new.json", "w") as f:
     json.dump(metric_dict, f, indent=4)
 

src/unit_exp.py

@@ -4,6 +4,8 @@
 from src.utils.model import *
 import os
 
+import numpy as np
+
 def set_seed(seed: int) -> None:
     """
     Sets the seed to make everything deterministic, for reproducibility of experiments
@@ -85,17 +87,17 @@ def train_single_model(param_dict: dict):
         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).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=np.sqrt(embd_dim), 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).to(device)
+        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, seq_len=input_token, seed=seed, use_dist_layer=True).to(device)
+        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=2, num_layers=2, n_dist=np.sqrt(embd_dim),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).to(device)
+        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:
         raise ValueError(f"Unknown model_id: {model_id}")
 

src/utils/driver.py

@@ -212,7 +212,7 @@ def __init__(self, shp, vocab_size, embd_dim, input_token=2, init_scale=1., weig
                 linear_list.append(DistLayer(shp[i], shp[i+1], n=n))
 
         self.embedding = nn.Embedding(vocab_size, embd_dim)
-        nn.init.normal_(self.embedding.weight, mean=0, std=1/np.sqrt(embd_dim))
+        nn.init.normal_(self.embedding.weight, mean=0, std=1/np.sqrt(embd_dim)*init_scale)
         #self.embedding = torch.nn.Parameter(torch.normal(0,1/torch.tensor(embd_dim),size=(vocab_size, embd_dim))*init_scale)
 #        self.embedding = torch.nn.Parameter(torch.normal(0,1,size=(vocab_size, embd_dim))*init_scale)
         self.linears = nn.ModuleList(linear_list)
@@ -256,15 +256,15 @@ def pred_logit(self, x):
 
 # 2-Layer Transformer Model with Explicit Residual Connections
 class ToyTransformer(customNNModule):
-    def __init__(self, vocab_size, d_model, nhead, num_layers, seq_len = 16, use_dist_layer = False, seed=0):
+    def __init__(self, vocab_size, d_model, nhead, num_layers, seq_len = 16, init_scale=1.,use_dist_layer = False, seed=0, n_dist=1.):
         super(ToyTransformer, self).__init__()
 
         torch.manual_seed(seed)
         np.random.seed(seed)
 
 
         self.embedding = nn.Embedding(vocab_size, d_model)
-        nn.init.normal_(self.embedding.weight, mean=0, std=1/np.sqrt(d_model))
+        nn.init.normal_(self.embedding.weight, mean=0, std=1/np.sqrt(d_model)*init_scale)
         self.positional_encoding = nn.Parameter(torch.randn(seq_len, d_model))
 
         # Define transformer encoder layers
@@ -275,7 +275,7 @@ def __init__(self, vocab_size, d_model, nhead, num_layers, seq_len = 16, use_dis
         ])
         self.use_dist_layer = use_dist_layer
         if use_dist_layer:
-            self.dist = DistLayer(d_model, vocab_size, n=1., eps=1e-4, bias=False)
+            self.dist = DistLayer(d_model, vocab_size, n=n_dist, eps=1e-4, bias=False)
         self.fc = nn.Linear(d_model, vocab_size)
         self.vocab_size = vocab_size