pytorch-cka
The fastest PyTorch library for layer-wise model similarity via Centered Kernel Alignment
Project Overview
pytorch-cka is a PyTorch library for measuring how similar the internal representations of two neural networks are, using Centered Kernel Alignment (CKA) from Kornblith et al. (ICML 2019). Built for scale, it uses vectorized operations and GPU acceleration to compute CKA across all 18 layers of ResNet-18 on CIFAR-10 up to 44x faster than existing implementations on an NVIDIA H100, while explicit memory deallocation keeps large comparisons within budget. It supports HuggingFace models, DataParallel, and DDP, ships heatmap and line-chart visualizations, and has surpassed 5K downloads on PyPI.