Update topk_sae.py

topk_sae.py

@@ -1,11 +1,9 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import torch.optim as optim
 import numpy as np
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm import tqdm
-import wandb
 import os
 import glob
 
@@ -153,109 +151,4 @@ def init_from_data_(ae, data_sample):
     # encoder as transpose of decoder
     ae.encoder.weight.data = ae.decoder.weight.t().clone()
 
-    nn.init.zeros_(ae.latent_bias)
-
-def train(ae, train_loader, optimizer, epochs, k, auxk_coef, multik_coef, clip_grad=None, save_dir="../models", model_name=""):
-    os.makedirs(save_dir, exist_ok=True)
-    step = 0
-    num_batches = len(train_loader)
-    for epoch in range(epochs):
-        ae.train()
-        total_loss = 0
-        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"):
-            optimizer.zero_grad()
-            x = batch[0].to(device)
-            recons, info = ae(x)
-            loss, recons_loss, auxk_loss = loss_fn(ae, x, recons, info, auxk_coef, multik_coef)
-            loss.backward()
-            step += 1
-
-            # calculate proportion of dead latents (not fired in last num_batches = 1 epoch)
-            dead_latents_prop = (ae.stats_last_nonzero > num_batches).float().mean().item()
-
-            wandb.log({
-                "total_loss": loss.item(),
-                "reconstruction_loss": recons_loss.item(),
-                "auxiliary_loss": auxk_loss.item(),
-                "dead_latents_proportion": dead_latents_prop,
-                "l0_norm": k,
-                "step": step
-            })
-            
-            unit_norm_decoder_grad_adjustment_(ae)
-            
-            if clip_grad is not None:
-                torch.nn.utils.clip_grad_norm_(ae.parameters(), clip_grad)
-            
-            optimizer.step()
-            unit_norm_decoder_(ae)
-            
-            total_loss += loss.item()
-        
-        avg_loss = total_loss / len(train_loader)
-        print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")
-
-        # Delete previous model saves for this configuration
-        for old_model in glob.glob(os.path.join(save_dir, f"{model_name}_epoch_*.pth")):
-            os.remove(old_model)
-
-        # Save new model
-        save_path = os.path.join(save_dir, f"{model_name}_epoch_{epoch+1}.pth")
-        torch.save(ae.state_dict(), save_path)
-        print(f"Model saved to {save_path}")
-
-def main():
-    d_model = 1536
-    n_dirs = 3072 #9216
-    k = 64 #64
-    auxk = k*2 #256
-    multik = 128
-    batch_size = 1024
-    lr = 1e-4
-    auxk_coef = 1/32
-    clip_grad = 1.0
-    multik_coef = 0 # turn it off
-
-    csLG = False
-
-    # Create model name
-    model_name = f"{k}_{n_dirs}_{auxk}_auxk" if not csLG else f"{k}_{n_dirs}_{auxk}_auxk_csLG"
-    epochs = 50 if not csLG else 137
-
-    wandb.init(project="saerch", name=model_name, config={
-        "n_dirs": n_dirs,
-        "d_model": d_model,
-        "k": k,
-        "auxk": auxk,
-        "batch_size": batch_size,
-        "lr": lr,
-        "epochs": epochs,
-        "auxk_coef": auxk_coef,
-        "multik_coef": multik_coef,
-        "clip_grad": clip_grad,
-        "device": device.type
-    })
-
-    if not csLG:
-        data = np.load("../data/vector_store_astroPH/abstract_embeddings.npy")
-        print("Doing astro.ph...")
-    else:
-        data = np.load("../data/vector_store_csLG/abstract_embeddings.npy")
-        print("Doing csLG...")
-    data_tensor = torch.from_numpy(data).float()
-    # Print shape
-    print(f"Data shape: {data_tensor.shape}")
-    dataset = TensorDataset(data_tensor)
-    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
-
-    ae = FastAutoencoder(n_dirs, d_model, k, auxk, multik).to(device)
-    init_from_data_(ae, data_tensor[:10000].to(device))
-
-    optimizer = optim.Adam(ae.parameters(), lr=lr)
-
-    train(ae, train_loader, optimizer, epochs, k, auxk_coef, multik_coef, clip_grad=clip_grad, model_name=model_name)
-
-    wandb.finish()
-
-if __name__ == "__main__":
-    main()
+    nn.init.zeros_(ae.latent_bias)