TopK_SAE.py

"""
Implementation of TopK Sparse Autoencoder with mixed precision training.
This autoencoder variant enforces sparsity by keeping only the k highest magnitude 
activations in the hidden layer.
"""

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import h5py
import wandb
from pathlib import Path
from tqdm import tqdm
from typing import Optional, Tuple, List
import matplotlib.pyplot as plt
import numpy as np
import time

class ActivationsDataset(Dataset):
    """Dataset for loading stored model activations."""
    def __init__(self, activations_path: str):
        try:
            self.file = h5py.File(activations_path, 'r')
        except Exception as e:
            raise RuntimeError(f"Failed to open HDF5 file at {activations_path}: {str(e)}")

        # Check if required datasets exist
        if 'activations' not in self.file:
            self.file.close()
            raise KeyError(
                f"HDF5 file at {activations_path} does not contain 'activations' dataset. "
                f"Available keys: {list(self.file.keys())}"
            )
        if 'paths' not in self.file:
            self.file.close()
            raise KeyError(
                f"HDF5 file at {activations_path} does not contain 'paths' dataset. "
                f"Available keys: {list(self.file.keys())}"
            )

        self.activations = self.file['activations']
        self.paths = self.file['paths']

        print(f"Loaded activations dataset with shape: {self.activations.shape}")
        print(f"Number of samples: {len(self.activations)}")

    def __len__(self) -> int:
        return len(self.activations)
        
    def __getitem__(self, idx: int) -> torch.Tensor:
        """Return a single item from the dataset."""
        # Convert to tensor and half precision, but don't pin memory here
        return torch.from_numpy(self.activations[idx]).half()
    
    def __del__(self):
        self.file.close()

class TopKSAE(nn.Module):
    """
    TopK Sparse Autoencoder implementation with mixed precision training.
    
    Args:
        input_dim: Dimension of input features
        expansion_factor: Factor to multiply input dimension by for hidden layer size
        k: Number of top activations to keep (others set to zero)
        learning_rate: Initial learning rate
        lr_scheduler: Learning rate scheduler type ('cosine', 'linear', or None)
        pretrained_path: Optional path to load pretrained weights from
    """
    def __init__(
        self,
        input_dim: int,
        expansion_factor: float,
        k: int,
        learning_rate: float,
        lr_scheduler: Optional[str] = None,
        pretrained_path: Optional[str] = None
    ):
        super().__init__()

        self.half()  # Use half precision for model weights
        
        self.input_dim = input_dim
        self.hidden_dim = int(input_dim * expansion_factor)
        self.k = k
        self.learning_rate = learning_rate
        self.lr_scheduler_type = lr_scheduler
        
        # Initialize encoder and decoder
        self.encoder = nn.Linear(input_dim, self.hidden_dim)
        self.decoder = nn.Linear(self.hidden_dim, input_dim)
        
        # Initialize with scaled random weights
        nn.init.kaiming_normal_(self.encoder.weight)
        nn.init.zeros_(self.encoder.bias)
        nn.init.kaiming_normal_(self.decoder.weight)
        nn.init.zeros_(self.decoder.bias)
        
        # Initialize mixed precision scaler
        self.scaler = torch.cuda.amp.GradScaler()
        
        # Load pretrained weights if provided
        if pretrained_path:
            self.load_state_dict(torch.load(pretrained_path))
            
    def top_k_activation(self, x: torch.Tensor) -> torch.Tensor:
        """Apply top-k sparsification to activations."""
        with torch.amp.autocast('cuda'):
            values, indices = torch.topk(x.abs(), self.k, dim=1)
            mask = torch.zeros_like(x)
            mask.scatter_(1, indices, 1)
            return x * mask
        
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward pass returning both hidden activations and reconstruction."""
        with torch.amp.autocast('cuda'):
            encoded = self.encoder(x)
            sparse_encoded = self.top_k_activation(encoded)
            decoded = self.decoder(sparse_encoded)
        return sparse_encoded, decoded
        
    def get_optimizer_and_scheduler(self, num_training_steps: int):
        """Setup optimizer and learning rate scheduler."""
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
        
        if self.lr_scheduler_type == 'cosine':
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizer, num_training_steps
            )
        elif self.lr_scheduler_type == 'linear':
            scheduler = torch.optim.lr_scheduler.LinearLR(
                optimizer,
                start_factor=1.0,
                end_factor=0.1,
                total_iters=num_training_steps
            )
        else:
            scheduler = None
            
        return optimizer, scheduler
    
    def dataloader(
        self,
        dataset_path: str,
        batch_size: int,
        num_workers: int,
        distributed: bool = False
    ) -> DataLoader:
        """Create dataloader for training/inference."""
        dataset = ActivationsDataset(dataset_path)

        if distributed:
            sampler = torch.utils.data.distributed.DistributedSampler(dataset)
        else:
            sampler = None

        return DataLoader(
            dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=True,  # Let the DataLoader handle pinning
            sampler=sampler,
            shuffle=sampler is None,
            drop_last=True,
            persistent_workers=True
        )
        
    def training_step(
        self,
        batch: torch.Tensor,
        optimizer: torch.optim.Optimizer
    ) -> torch.Tensor:
        """Single training step with mixed precision."""
        # start_time = time.time()

        optimizer.zero_grad()
        # zero_grad_time = time.time()

        # Compute loss in mixed precision
        with torch.amp.autocast('cuda'):
            _, reconstruction = self(batch)
            loss = F.mse_loss(reconstruction, batch)
        # forward_time = time.time()

        # Scale loss and backpropagate
        self.scaler.scale(loss).backward()
        # backward_time = time.time()

        self.scaler.step(optimizer)
        self.scaler.update()
        # step_time = time.time()

        # print(f"Zero grad: {zero_grad_time - start_time:.4f}s")
        # print(f"Forward: {forward_time - zero_grad_time:.4f}s")
        # print(f"Backward: {backward_time - forward_time:.4f}s")
        # print(f"Optimizer step: {step_time - backward_time:.4f}s")
        # print(f"Total step: {step_time - start_time:.4f}s")

        return loss
    
    def get_latents(self, images: torch.Tensor) -> torch.Tensor:
        """Get latent representations for input images."""
        self.eval()
        with torch.no_grad():
            with torch.amp.autocast('cuda'):  # Use mixed precision for inference too
                latents, _ = self(images)
        return latents
        
    def get_top_images(
        self,
        neuron_number: int,
        dataset_path: str,
        output_path: str,
        batch_size: int = 32,
        split: str = 'val',
        num_workers: int = 8,
        device: str = 'cuda',
        top_k: int = 9,
        dataset_name: str = 'imagenet',
        model_library: str = 'clip',
        model_version: str = 'ViT-B/32',
        block_index: int = 5,
        block_element: str = 'mlp',
        expansion_factor: float = 8.0,
        token_to_save: str = 'cls'
    ):
        """Find and save images that most activate a specific neuron."""
        self.eval()
        self.to(device)

        # Setup dataloader with shuffle=False to maintain consistent ordering

        # Create dataset
        dataset = ActivationsDataset(dataset_path)
    
        # Create dataloader with no shuffling
        dataloader = DataLoader(
            dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=True,
            shuffle=False  # Explicit no shuffling
        )

        # Track all activations with their paths
        activations_with_paths = []

        print(f"Processing activations for neuron {neuron_number}")
        with torch.no_grad():
            for batch in tqdm(dataloader):
                batch = batch.to(device)
                latents = self.get_latents(batch)
                activations = latents[:, neuron_number].cpu().numpy()

                # Get paths for this batch from the dataset
                paths = [str(p).strip("b'\"") for p in batch[1]] if isinstance(batch, tuple) else \
                       [str(p).strip("b'\"") for p in dataloader.dataset.paths[len(activations_with_paths):len(activations_with_paths) + len(batch)]]

                # Store all activations with their paths
                activations_with_paths.extend(zip(activations, paths))

        print(f"Found {len(activations_with_paths)} total activations")

        # Sort all activations and take top k
        activations_with_paths.sort(key=lambda x: x[0], reverse=True)
        top_k_pairs = activations_with_paths[:top_k]

        if not top_k_pairs:
            print(f"No activations found for neuron {neuron_number}")
            return

        top_activations, top_image_paths = zip(*top_k_pairs)

        if top_activations[0] > 0:
            # Create output directory
            output_dir = Path(os.path.join(output_path, dataset_name, split, model_library, 
                                         model_version, token_to_save, f'lr_{self.learning_rate}', f'k_{self.k}', f'ef_{expansion_factor}', str(block_index), block_element))
            output_dir.mkdir(parents=True, exist_ok=True)

            # Save information about the activations
            info_path = output_dir / f'neuron_{neuron_number}_info.txt'
            with open(info_path, 'w') as f:
                f.write(f"Top {top_k} activations for neuron {neuron_number}\n")
                f.write("-" * 50 + "\n")
                for act, path in zip(top_activations, top_image_paths):
                    f.write(f"Activation: {act:.6f}\n")
                    f.write(f"Image path: {path}\n")
                    f.write("-" * 50 + "\n")

            # Create visualization
            rows = int(np.ceil(np.sqrt(top_k)))
            cols = int(np.ceil(top_k / rows))
            fig, axes = plt.subplots(rows, cols, figsize=(4*cols, 4*rows))

            if rows == 1 and cols == 1:
                axes = np.array([[axes]])
            elif rows == 1:
                axes = axes[np.newaxis, :]
            elif cols == 1:
                axes = axes[:, np.newaxis]

            # Plot images
            for idx, (act, img_path) in enumerate(zip(top_activations, top_image_paths)):
                if idx >= top_k:
                    break

                row, col = idx // cols, idx % cols
                try:
                    img = plt.imread(img_path)
                    axes[row, col].imshow(img)
                    axes[row, col].axis('off')
                    axes[row, col].set_title(f'Activation: {act:.3f}')
                except Exception as e:
                    print(f"Error loading image {img_path}: {str(e)}")
                    axes[row, col].text(0.5, 0.5, 'Image load error', ha='center', va='center')
                    axes[row, col].axis('off')

            # Clear any unused subplots
            for idx in range(len(top_activations), rows * cols):
                row, col = idx // cols, idx % cols
                axes[row, col].axis('off')

            plt.tight_layout()
            fig.savefig(output_dir / f'neuron_{neuron_number}_top_activations.pdf', bbox_inches='tight')
            fig.savefig(output_dir / f'neuron_{neuron_number}_top_activations.png', bbox_inches='tight', dpi=300)
            plt.close()

            print(f"Plots and info saved to: {output_dir}")
        else:
            print(f"No activations found for neuron {neuron_number}")