Sparse Autoencoder

# First, download CelebA dataset.
import os
import getpass

os.environ['KAGGLE_USERNAME'] = getpass.getpass("Username")
os.environ['KAGGLE_KEY'] = getpass.getpass("API Key")

!kaggle datasets download -d jessicali9530/celeba-dataset -p CelebA

# Our images unzip to `./img_align_celeba/img_align_celeba` (relative to the current directory; not the CelebA folder for some reason.)
!unzip CelebA/celeba-dataset.zip -d data > /dev/null

import matplotlib.pyplot as plt
import torch
import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

def visualize_4x4(images: torch.Tensor, title=None):
    # Shape: [N, C, H, W]

    # Make a grid of images
    grid = torchvision.utils.make_grid(images, nrow=4).numpy()

    # Render
    plt.figure(figsize=(4, 4))
    plt.title(title or "Celebrities!")
    plt.imshow(grid.transpose((1, 2, 0)))
    plt.axis("off")
    plt.show()

# Now, train a contrastive model. Use CelebA dataset.
# dataset[i] = (PIL.Image.Image, <dummy class label>)
dataset = ImageFolder(root="./data", transform=T.Compose([
    T.ToTensor(),
    T.Resize((224, 224)),
]))

visualize_4x4([dataset[i][0] for i in range(16)])

from torchvision.models import resnet

# Decide whether to use the GPU. This speeds up training a *LOT* if it's available.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# Remove the fully-connected layer because we're interested in embeddings.
model.fc = torch.nn.Identity()
model = model.to(device)

import tqdm

batch_size = 128
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

optim = torch.optim.Adam(model.parameters(), lr=1e-3)

N_EPOCHS = 1
for epoch in range(N_EPOCHS):
  with tqdm.tqdm(dataloader, desc="Training...") as pbar:
    for (images, _) in pbar:
      images = images.to(device)
      # Train model to disambiguate between these images.
      # Use embeddings as features. Instead of retraining
      # all of EfficientNet, add a linear layer at the end
      # to get embeddings. Rotate vector around first dimension.
      # 180º rotated images should have negative embeddings of each other.
      images_rotated = torch.rot90(images, 2, (2, 3))
      # [N, embedding_dim]
      embeddings = model(images)
      embeddings_rotated = model(images_rotated)

      # ORIENTATION LOSS
      loss_orientation = torch.nn.functional.mse_loss(embeddings, -embeddings_rotated)

      # CONTRASTIVE LOSS
      # Dot each of these embeddings with each other and treat these as logits.
      # Then use Cross Entropy loss.
      embeddings_normalized = torch.nn.functional.normalize(embeddings, dim=1)
      scores = embeddings @ embeddings.T
      loss_contrastive = torch.nn.functional.cross_entropy(scores, torch.arange(len(scores), device=device))

      loss = loss_contrastive + loss_orientation

      optim.zero_grad()
      loss.backward()
      optim.step()
      pbar.set_postfix({"loss": loss.item(), "contrastive_loss": loss_contrastive.item(), "orientation_loss": loss_orientation.item()})

torch.save(model.state_dict(), "resnet_18_celebface_state_dict.pt")

results = []
model.eval()

# Notably, we *DO NOT SHUFFLE*!
dataloader_eval = DataLoader(dataset, batch_size=batch_size, shuffle=False)

with torch.no_grad():
  with tqdm.tqdm(dataloader_eval, desc="Generating output vectors...") as pbar:
    for (images, _) in pbar:
      images = images.to(device)
      embeddings = model(images)
      results.append(embeddings.to('cpu'))

# Took about 17m50s for 1583 batches (202599 input faces)

# Generate a single tensor
face_vectors = torch.cat(results, dim=0)

# See how much data we got :0
# Each vector has 512 dimensions.
print(face_vectors.shape)

# Store this tensor
torch.save(face_vectors, "face_vectors.pt")

class SparseAutoencoder(torch.nn.Module):
  def __init__(self, input_feature_dim, sparse_feature_dim):
    super().__init__() # necessary for any torch.nn.Module subclass

    self.weight_encoder = torch.nn.Parameter(torch.zeros((input_feature_dim, sparse_feature_dim), dtype=torch.float32))
    self.weight_decoder = torch.nn.Parameter(torch.zeros((sparse_feature_dim, input_feature_dim), dtype=torch.float32))
    self.bias_decoder = torch.nn.Parameter(torch.zeros(input_feature_dim, dtype=torch.float32))
    self.bias_encoder = torch.nn.Parameter(torch.zeros(sparse_feature_dim, dtype=torch.float32))
    # initialize weight matrices
    torch.nn.init.normal_(self.weight_encoder)
    torch.nn.init.normal_(self.weight_decoder)

  def encode(self, input_features):
    # (n, input_feature_dim)
    xbar = input_features - self.bias_decoder
    # (n, sparse_feature_dim)
    f = torch.nn.functional.relu(xbar @ self.weight_encoder + self.bias_encoder)
    return f

  def decode(self, sparse_features):
    # (n, sparse_feature_dim)
    xhat = sparse_features @ self.weight_decoder + self.bias_decoder
    return xhat

def sparse_autoencoder_loss(model, input_features, l1_lambda):
  # Assume input is batched: (n, input_feature_dim)
  f = model.encode(input_features)
  x_hat = model.decode(f)
  # L_reconstruction = torch.nn.functional.mse_loss(input_features, x_hat)
  L_reconstruction = torch.norm(input_features - x_hat, 2)
  L_complexity = l1_lambda * torch.norm(f, 1)
  return L_reconstruction, L_complexity

import PIL.Image

import torch

face_vectors = torch.load("face_vectors.pt")

sparse_feature_dim = 512
sae = SparseAutoencoder(input_feature_dim=512, sparse_feature_dim=sparse_feature_dim).to('cuda')
sae_optim = torch.optim.Adam(sae.parameters(), lr=1e-5, weight_decay=1e-4)

import tqdm

# l_recon = 30 or so.
# Takes like 1 minute train this many epochs, because of CUDA and the small size of the dataset.
for epoch in range(30):
  i = 0
  sae_batch_size = 512
  N = len(face_vectors)
  order = torch.randperm(N)
  with tqdm.tqdm(total=N, desc="Training sparse autoencoder...") as pbar:
    while i < N:
      batch = face_vectors[order[i:i + 512]].to('cuda')
      # L_reconstruction, L_complexity = sparse_autoencoder_loss(sae, torch.nn.functional.normalize(batch) * torch.sqrt(torch.tensor(512.0)), l1_lambda=0.1)
      L_reconstruction, L_complexity = sparse_autoencoder_loss(sae, torch.nn.functional.normalize(batch), l1_lambda=0.1)
      loss_sae = L_reconstruction + L_complexity
      sae_optim.zero_grad()
      loss_sae.backward()
      sae_optim.step()
      pbar.set_postfix({"loss_sae": loss_sae.item(), "recon": L_reconstruction.item(), "compl": L_complexity.item()})
      pbar.update(sae_batch_size)
      i += sae_batch_size

# Save our sparse autoencoder.
torch.save(sae.state_dict(), "sparse_autoencoder_weights.pt")

# Let's sort images by their activation of certain neurons.
sparse_features = []

with torch.no_grad():
  i = 0
  sae_batch_size = 512
  N = len(face_vectors)
  with tqdm.tqdm(total=N, desc="Extracting sparse features...") as pbar:
    while i < N:
      batch = face_vectors[i:i + 512].to('cuda')
      sparse_features.append(sae.encode(batch).to('cpu'))
      i += sae_batch_size
      pbar.update(sae_batch_size)

feats = torch.cat(sparse_features, dim=0)

# Save our sparse features.
torch.save(feats, "face_vectors_sparse.pt")

feats_dense_normed = torch.nn.functional.normalize(face_vectors)

for feature_id in range(16):
  # Sort images by feature `feature_id`.
  order = torch.argsort(feats_dense_normed[:, feature_id], descending=True)
  visualize_4x4(
      [dataset[order[i]][0] for i in range(16)],
      title="Maximally-Activating Examples",
  )

feats_normed = torch.nn.functional.normalize(feats)

for feature_id in range(16):
  # Sort images by feature `feature_id`.
  order = torch.argsort(feats_normed[:, feature_id], descending=True)
  visualize_4x4(
      [dataset[order[i]][0] for i in range(16)],
      title="Maximally-Activating Examples",
  )

import random
import time

test_features = [*range(16, 32)]
orders = []

for feature_id in test_features:
  # Sort images by feature `feature_id`.
  order_dense = torch.argsort(feats_dense_normed[:, feature_id], descending=True)
  order_sparse = torch.argsort(feats_normed[:, feature_id], descending=True)
  orders.append((order_dense, "dense:" + str(feature_id)))
  orders.append((order_sparse, "sparse:" + str(feature_id)))

random.shuffle(orders)
results = []
for (order, feature_id) in orders:
  visualize_4x4(
      [dataset[order[i]][0] for i in range(16)],
      title="Maximally-Activating Examples",
  )
  time.sleep(0.5)
  rating = int(input("How visually-consistent do images for this feature appear? (1 - 7)\n"))
  results.append((feature_id, rating))

# Let's see how the ratings turned out.

ratings = {"sparse": [], "dense": []}
for (feature_id, rating) in results:
  vector_type = feature_id.split(":")[0]
  ratings[vector_type].append(rating)

# Create a figure and axis
fig, ax = plt.subplots()

# Create histograms
n, bins, patches = ax.hist([ratings['sparse'], ratings['dense']], bins=[1, 2, 3, 4, 5, 6, 7], label=["Sparse Features", "Dense Features"], alpha=0.5)

# Calculate the positions for centered ticks
tick_positions = 0.5 * (bins[:-1] + bins[1:])
tick_labels = [str(int(tick)) for tick in tick_positions]

# Set the x-tick positions and labels
ax.set_xticks(tick_positions)
ax.set_xticklabels(tick_labels)

# Set labels and legend
ax.set_xlabel("Interpretability Rating 1-7")
ax.set_title("Human-Annotated Blind Interpretability Ratings")
ax.legend()

# Show the plot
plt.show()

# Let's get the mean and standard deviation of these ratings through a Pandas summary.
# We'll also do a t-test to see if sparse features are truly more explainable than dense features.
import pandas as pd
from scipy.stats import ttest_ind

display(pd.DataFrame(ratings).describe())

ttest_result = ttest_ind(ratings['sparse'], ratings['dense'])
print(f"t-stat: {ttest_result.statistic:.4f}, pvalue: {ttest_result.pvalue:.4f}")

visualize_4x4([dataset[267][0]] * 16)

import PIL.Image
PIL.Image.open("data/img_align_celeba/img_align_celeba/000268.jpg")