compute_cka.py

from modelzoo import xception, separable_net, gabor_pyramid, dorsalnet, decoder
from loaders import airsim
from models import extract_subnet_dict

import argparse
import datetime
import itertools
import os
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse

import numpy as np
import torch
from torch import nn
from torch import optim

import torch.autograd.profiler as profiler
from torchvision import transforms
import torchvision.models as models
import torch.nn.functional as F

import models

from transforms import ThreedGaussianBlur, ThreedExposure

import train_heading

import wandb

from models import preprocess_data
from research_code.cka_step4 import cka, multi_cka

from models import get_aggregator


def main(args):
    print("Main")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device == "cpu":
        print("No CUDA! Sad!")

    (
        trainset,
        tuneset,
        reportset,
        train_transform,
        eval_transform,
        start_sz,
    ) = train_heading.get_dataset(args)

    reportloader = torch.utils.data.DataLoader(
        reportset, batch_size=args.batch_size, shuffle=True, pin_memory=True
    )

    print("Init models")

    subnet, activations, metadata = train_heading.get_subnet(args, start_sz)
    aggregator = get_aggregator(metadata, args)

    subnet.to(device=device)
    subnet.eval()

    # Use the report fold as a means of evaluating the CKA alignments.
    layer_responses = []
    for layer_num, layer_name in enumerate(metadata["layers"].keys()):
        args.layer = layer_num  # For backwards compatibility
        args.layer_name = layer_name
        args.subset = 0
        X, Y = preprocess_data(
            reportloader, subnet, aggregator, activations, metadata, args
        )

        # Pick a time and position slice in the middle.
        X = X.reshape((X.shape[0], -1, 4, args.aggregator_sz, args.aggregator_sz))
        X = X[:, :, 2, args.aggregator_sz // 2, args.aggregator_sz // 2].reshape(
            (X.shape[0], -1)
        )
        layer_responses.append(X)

    alignments = multi_cka(layer_responses)
    np.save("cka.npy", alignments)


if __name__ == "__main__":
    desc = "Measure CKA between layers of a deep neural net"
    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument(
        "--features",
        default="gaborpyramid3d",
        type=str,
        help="Sub-model type (currently, either xception2d, gaborpyramid2d, gaborpyramid3d",
    )
    parser.add_argument(
        "--dataset", default="airsim", help="Dataset (currently airsim only)"
    )
    parser.add_argument("--data_root", default="./data_derived", help="Data path")
    parser.add_argument("--ckpt_root", default="./checkpoints", help="Data path")
    parser.add_argument(
        "--output_dir", default="./models", help="Output path for models"
    )
    parser.add_argument(
        "--softmax",
        default=False,
        help="Use softmax objective rather than regression",
        action="store_true",
    )
    parser.add_argument(
        "--slowfast_root", default="", help="Path where SlowFast is installed"
    )
    parser.add_argument(
        "--subsample_layers",
        default=False,
        help="Subsample layers (saves disk space & mem)",
        action="store_true",
    )
    parser.add_argument(
        "--cache_root", default="./cache", help="Precomputed cache path"
    )
    parser.add_argument("--batch_size", default=4, type=int, help="Batch size")
    parser.add_argument(
        "--aggregator",
        default="average",
        type=str,
        help="What kind of aggregator to use",
    )
    parser.add_argument(
        "--aggregator_sz",
        default=8,
        type=int,
        help="The size the aggregator will be used with",
    )

    args = parser.parse_args()
    main(args)