44x faster CKA computation across 18 representational layers of ResNet-18 models on CIFAR-10 using NVIDIA H100 GPUs

• ⚡️ Fastest among CKA libraries thanks to vectorized ops & GPU acceleration
• 📦 Efficient memory management with explicit deallocation
• 🧠 Supports HuggingFace models, DataParallel, and DDP
• 🎨 Customizable visualizations: heatmaps and line charts

📦 Installation

Requires Python 3.10+

# Using pip
pip install pytorch-cka

# Using uv
uv add pytorch-cka

👟 Quick Start

Basic Usage

from cka import compute_cka
from torch.utils.data import DataLoader
from torchvision.models import resnet18, resnet34

resnet_18 = resnet18(pretrained=True)
resnet_34 = resnet34(pretrained=True)

dataloader1 = Dataloader(your_dataset1, batch_size=bach_size, shuffle=False, num_workers=4)
dataloader2 = Dataloader(your_dataset2, batch_size=bach_size, shuffle=False, num_workers=4)
dataloader3 = Dataloader(your_dataset3, batch_size=bach_size, shuffle=False, num_workers=4)
dataloaders = [dataloader1, dataloader2, dataloader3]

layers = [
    'conv1',
    'layer1.0.conv1',
    'layer2.0.conv1',
    'layer3.0.conv1',
    'layer4.0.conv1',
    'fc',
]

cka_matrices = compute_cka(
    resnet_18,
    resnet_34,
    dataloaders,
    layers=layers,
    device=device,
)

for cka_matrix in cka_matrices:
    print(cka_matrix)

From Pre-extracted Features

If you already have feature matrices, compute CKA without models or dataloaders:

from cka import cka_from_features

# Single layer: (n_samples, feature_dim)
cka_matrix = cka_from_features(features_x, features_y)

# Multi-layer: (n_layers, n_samples, feature_dim)
cka_matrix = cka_from_features(multi_layer_x, multi_layer_y)

# Varying feature dims: list of 2D tensors
cka_matrix = cka_from_features(
    [layer1_x, layer2_x],
    [layer1_y, layer2_y, layer3_y],
)

Visualization

Heatmap

from cka import plot_cka_heatmap

fig, ax = plot_cka_heatmap(
    cka_matrix,
    layers1=layers,
    layers2=layers,
    model1_name="ResNet-18 (pretrained)",
    model2_name="ResNet-18 (random init)",
    annot=False,          # Show values in cells
    cmap="inferno",       # Colormap
)

Self-comparison	Cross-model

Trend Plot

from cka import plot_cka_trend

# Plot diagonal (self-similarity across layers)
diagonal = torch.diag(matrix)

fig, ax = plot_cka_trend(
    layer_trends,
    x_values=epochs,
    labels=RESNET18_LAYERS,
    markers=['o'],
    xlabel='Epoch',
    ylabel='CKA Score',
    title='Pretrained vs. Fine-tuned Across Epochs (ResNet-18)',
    legend=True,
)

fig, ax = plot_cka_layer_trend(
    cka_matrices,
    layers=RESNET18_LAYERS,
    labels=cka_loader_names,
    ylabel='CKA Score',
    title='Pretrained vs. Fine-tuned Across Layers (ResNet-18)',
    legend=True,
)

CKA Score Trend Across Epochs	CKA Score Trend Across Layers

📚 References

Kornblith, Simon, et al. "Similarity of Neural Network Representations Revisited." ICML 2019.