A bit of cleanup, file naming clarity

Author: thedch

autoencoder/autoencoder.py autoencoder/autoencoder_architecture.pyautoencoder/autoencoder.py renamed to autoencoder/autoencoder_architecture.py

@@ -58,7 +58,7 @@ def load_text_data(self):
     def load_transformer_model(self):
         """Loads the GPT model with pre-trained weights."""
         ckpt_path = os.path.join(self.base_dir, 'transformer', self.gpt_ckpt_dir, 'ckpt.pt')
-        checkpoint = torch.load(ckpt_path, map_location=self.device)
+        checkpoint = torch.load(ckpt_path, map_location=self.device, weights_only=False)
         gpt_conf = GPTConfig(**checkpoint['model_args'])
         transformer = HookedGPT(gpt_conf)
         state_dict = checkpoint['model']
@@ -100,7 +100,7 @@ def load_autoencoder_model(self):
         state_dict = autoencoder_ckpt['autoencoder']
         n_features, n_ffwd = state_dict['encoder.weight'].shape  # H, F
         l1_coeff = autoencoder_ckpt['config']['l1_coeff']
-        from autoencoder import AutoEncoder
+        from autoencoder_architecture import AutoEncoder
 
         autoencoder = AutoEncoder(n_ffwd, n_features, lam=l1_coeff).to(self.device)
         autoencoder.load_state_dict(state_dict)
@@ -114,7 +114,7 @@ def get_text_batch(self, num_contexts):
         Y = torch.stack([torch.from_numpy(self.text_data[i + 1 : i + 1 + block_size].astype(np.int64)) for i in ix])
         return X.to(device=self.device), Y.to(device=self.device)
 
-    def get_autoencoder_data_batch(self, step, batch_size=8192):
+    def get_autoencoder_data_batch(self, step, batch_size: int):
         """
         Retrieves a batch of autoencoder data based on the step and batch size.
         It loads the next data partition if the batch exceeds the current partition.
@@ -141,14 +141,12 @@ def get_autoencoder_data_batch(self, step, batch_size=8192):
         return batch.to(self.device)
 
     def load_next_autoencoder_partition(self, partition_id):
-        """
-        Loads the specified partition of the autoencoder data.
-        """
+        """Loads the specified partition of the autoencoder data."""
         file_path = os.path.join(self.autoencoder_data_dir, f'{partition_id}.pt')
-        self.autoencoder_data = torch.load(file_path)
+        self.autoencoder_data = torch.load(file_path, weights_only=False)
         return self.autoencoder_data
 
-    def select_resampling_data(self, size=819200):
+    def select_resampling_data(self, size: int):
         """
         Selects a subset of autoencoder data for resampling, distributed evenly across partitions.
         """

autoencoder/prepare.py …toencoder/prepare_autoencoder_dataset.pyautoencoder/prepare.py renamed to autoencoder/prepare_autoencoder_dataset.py autoencoder/prepare.py renamed to autoencoder/prepare_autoencoder_dataset.py

@@ -2,14 +2,14 @@
 Train a Sparse AutoEncoder model
 
 Run on a macbook on a Shakespeare dataset as
-python train.py --dataset=shakespeare_char --gpt_ckpt_dir=out_sc_1_2_32 --eval_iters=1 --eval_batch_size=16 --batch_size=128 --device=cpu --eval_interval=100 --n_features=1024 --resampling_interval=150 --wandb_log=True
+python train.py --dataset=shakespeare_char --gpt_ckpt_dir=out_sc_1_2_32 --eval_iters=1 --eval_batch_size=16 --batch_size=128 --device=cpu --eval_interval=100 --n_features=1024 --resampling_interval=150
 """
 
 import os
 import torch
 import numpy as np
 import time
-from autoencoder import AutoEncoder
+from autoencoder_architecture import AutoEncoder
 from resource_loader import ResourceLoader
 from utils.plotting_utils import make_density_histogram
 
@@ -18,7 +18,7 @@
 dataset = 'openwebtext'
 gpt_ckpt_dir = 'out'
 # training
-n_features = 4096
+n_features = 4096  # aka n_latents
 batch_size = 8192  # batch size for autoencoder training
 l1_coeff = 3e-3
 learning_rate = 3e-4
@@ -33,8 +33,6 @@
 save_checkpoint = True  # whether to save model, optimizer, etc or not
 save_interval = 10000  # number of training steps after which a checkpoint will be saved
 out_dir = 'out'  # directory containing trained autoencoder model weights
-# wandb logging
-wandb_log = True
 # system
 device = 'cuda'
 # reproducibility
@@ -57,20 +55,15 @@
 
 gpt = resourceloader.transformer  # TODO: either it should be called transformer or gpt
 autoencoder = AutoEncoder(
-    n_inputs=(4 * resourceloader.transformer.config.n_embd),  # ?? why 4x?
+    n_inputs=(4 * resourceloader.transformer.config.n_embd),
     n_latents=n_features,
     lam=l1_coeff,
     resampling_interval=resampling_interval
 ).to(device)
 optimizer = torch.optim.Adam(autoencoder.parameters(), lr=learning_rate)
 
 ## prepare for logging and saving checkpoints
-run_name = f'{time.time():.2f}'
-if wandb_log:
-    raise DeprecationWarning('wandb is deprecated')
-    import wandb
-
-    wandb.init(project=f'sparse-autoencoder-{dataset}', name=run_name, config=config)
+run_name = time.strftime('%Y-%m-%d-%H%M')
 
 if save_checkpoint:
     ckpt_path = os.path.join(out_dir, dataset, run_name)
@@ -181,9 +174,6 @@
                 'feature_density/num_alive_neurons': len(log_feat_acts_density),
             }
         )
-        if wandb_log:
-            log_dict.update({'feature_density/feature_density_histograms': wandb.Image(feat_density_historgram)})
-            wandb.log(log_dict)
 
         autoencoder.train()
         print(f'Exiting evaluation mode at step = {step}')
@@ -198,7 +188,4 @@
             'feature_activation_counts': feat_acts_count,  # may be used later to identify alive vs dead neurons
         }
         print(f"saving checkpoint to {ckpt_path} at training step = {step}")
-        torch.save(checkpoint, os.path.join(ckpt_path, 'ckpt.pt'))
-
-if wandb_log:
-    wandb.finish()
+        torch.save(checkpoint, os.path.join(ckpt_path, 'ckpt.pt'))