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)