Minor tweaks

Author: david-baek

scripts/u_circle.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
 #SBATCH -t 16:00:00
-#SBATCH --gres=gpu:1
+#SBATCH --gres=gpu:a100: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
-
+for ARG in $(python -c "import numpy as np; print(' '.join(map(str, np.linspace(49, 100, 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

src/run_exp.py

@@ -19,7 +19,7 @@
 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=29, help='random seed')
+    parser.add_argument('--seed', type=int, default=49, 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')
 
@@ -74,7 +74,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 = [400, 500, 600, 700, 800, 900, 1000, 1100, 1200]
+data_size_list = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]
 for i in tqdm(range(len(data_size_list))):
     data_size = data_size_list[i]
     param_dict = {

src/unit_exp.py

@@ -51,7 +51,7 @@
 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}.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)
@@ -71,7 +71,7 @@
 elif data_id == "equivalence":
     aux_info["mod"] = 5
 elif data_id == "circle":
-    aux_info["p"] = 31
+    aux_info["p"] = 17
 else:
     raise ValueError(f"Unknown data_id: {data_id}")
 

src/utils/driver.py

@@ -91,21 +91,21 @@ def train_single_model(param_dict: dict):
         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)
     elif model_id == "H_transformer":
-        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=8, num_layers=1, seq_len=input_token, use_dist_layer=True).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=True).to(device)
     elif model_id == "standard_transformer":
-        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=8, num_layers=1, seq_len=input_token, 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).to(device)
     else:
         raise ValueError(f"Unknown model_id: {model_id}")
 
     # define dataloader
-    batch_size = 16
+    batch_size = 32
     train_dataset = ToyDataset(dataset['train_data_id'], dataset['train_label'])
     test_dataset = ToyDataset(dataset['test_data_id'], dataset['test_label'])
     train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
     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': 4000, 'learning_rate': 0.001, 'train_dataloader': train_dataloader, 'test_dataloader': test_dataloader, 'device': device})
+    ret_dic["results"] = model.train(param_dict={'num_epochs': 7000, 'learning_rate': 0.002, 'train_dataloader': train_dataloader, 'test_dataloader': test_dataloader, 'device': device})
     ret_dic["model"] = model
     ret_dic["dataset"] = dataset
 

src/utils/model.py

@@ -5,6 +5,8 @@
 import numpy as np
 import math
 
+import sys
+
 from tqdm import tqdm
 
 class customNNModule(nn.Module):
@@ -34,7 +36,7 @@ def train(self, param_dict: dict):
         counter = 0 
 
         optimizer = optim.AdamW(self.parameters(), lr=learning_rate, weight_decay=0.01)
-        lamb_reg = 0.1
+        lamb_reg = 0.01
         for epoch in tqdm(range(num_epochs)):
             train_loss = 0
             train_correct = 0
@@ -86,6 +88,7 @@ def train(self, param_dict: dict):
 
             if (epoch + 1) % 50 == 0 and verbose:
                 print(f"Epoch {epoch + 1}/{num_epochs}, Train Loss: {train_loss / len(train_dataloader):.4f}, Train Acc: {train_correct / train_total:.4f}, Test Loss: {test_loss / len(test_dataloader):.4f}, Test Acc: {test_correct / test_total:.4f}")
+                sys.stdout.flush()
 
             train_losses.append(train_loss / len(train_dataloader))
             test_losses.append(test_loss / len(test_dataloader))
@@ -252,16 +255,21 @@ 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):
+    def __init__(self, vocab_size, d_model, nhead, num_layers, seq_len = 16, use_dist_layer = False, seed=0):
         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))
         self.positional_encoding = nn.Parameter(torch.randn(seq_len, d_model))
 
         # Define transformer encoder layers
         self.layers = nn.ModuleList([
             nn.TransformerEncoderLayer(
-                d_model=d_model, nhead=nhead, dim_feedforward=64, batch_first=True
+                d_model=d_model, nhead=nhead, dim_feedforward=d_model*4, batch_first=True
             ) for _ in range(num_layers)
         ])
         self.use_dist_layer = use_dist_layer