similarity.py

import math

import torch
from tqdm import tqdm


def cos_similarity_cubed(clip_feats, target_feats, device='cuda', batch_size=10000, min_norm=1e-3):
    """
    Substract mean from each vector, then raises to third power and compares cos similarity
    Does not modify any tensors in place
    """
    with torch.no_grad():
        torch.cuda.empty_cache()
        
        clip_feats = clip_feats - torch.mean(clip_feats, dim=0, keepdim=True)
        target_feats = target_feats - torch.mean(target_feats, dim=0, keepdim=True)
        
        clip_feats = clip_feats**3
        target_feats = target_feats**3
        
        clip_feats = clip_feats/torch.clip(torch.norm(clip_feats, p=2, dim=0, keepdim=True), min_norm)
        target_feats = target_feats/torch.clip(torch.norm(target_feats, p=2, dim=0, keepdim=True), min_norm)
        
        similarities = []
        for t_i in range(math.ceil(target_feats.shape[1]/batch_size)):
            curr_similarities = []
            curr_target = target_feats[:, t_i*batch_size:(t_i+1)*batch_size].to(device).T
            for c_i in range(math.ceil(clip_feats.shape[1]/batch_size)):
                curr_similarities.append(curr_target @ clip_feats[:, c_i*batch_size:(c_i+1)*batch_size].to(device))
            similarities.append(torch.cat(curr_similarities, dim=1))
    return torch.cat(similarities, dim=0)

def cos_similarity(clip_feats, target_feats, device='cuda'):
    with torch.no_grad():
        clip_feats = clip_feats / torch.norm(clip_feats, p=2, dim=0, keepdim=True)
        target_feats = target_feats / torch.norm(target_feats, p=2, dim=0, keepdim=True)
        
        batch_size = 10000
        
        similarities = []
        for t_i in range(math.ceil(target_feats.shape[1]/batch_size)):
            curr_similarities = []
            curr_target = target_feats[:, t_i*batch_size:(t_i+1)*batch_size].to(device).T
            for c_i in range(math.ceil(clip_feats.shape[1]/batch_size)):
                curr_similarities.append(curr_target @ clip_feats[:, c_i*batch_size:(c_i+1)*batch_size].to(device))
            similarities.append(torch.cat(curr_similarities, dim=1))
    return torch.cat(similarities, dim=0)

def soft_wpmi(clip_feats, target_feats, top_k=100, a=10, lam=1, device='cuda',
                        min_prob=1e-7, p_start=0.998, p_end=0.97):
    top_k = min(top_k, target_feats.size(0))
    with torch.no_grad():
        torch.cuda.empty_cache()
        clip_feats = torch.nn.functional.softmax(a*clip_feats, dim=1)

        inds = torch.topk(target_feats, dim=0, k=top_k)[1]
        prob_d_given_e = []

        p_in_examples = p_start-(torch.arange(start=0, end=top_k)/top_k*(p_start-p_end)).unsqueeze(1).to(device)
        for orig_id in range(target_feats.shape[1]):
            
            curr_clip_feats = clip_feats.gather(0, inds[:,orig_id:orig_id+1].expand(-1,clip_feats.shape[1])).to(device)
            
            curr_p_d_given_e = 1+p_in_examples*(curr_clip_feats-1)
            curr_p_d_given_e = torch.sum(torch.log(curr_p_d_given_e+min_prob), dim=0, keepdim=True)
            prob_d_given_e.append(curr_p_d_given_e)
            torch.cuda.empty_cache()

        prob_d_given_e = torch.cat(prob_d_given_e, dim=0)
        #logsumexp trick to avoid underflow
        prob_d = (torch.logsumexp(prob_d_given_e, dim=0, keepdim=True) - 
                  torch.log(prob_d_given_e.shape[0]*torch.ones([1]).to(device)))
        mutual_info = prob_d_given_e - lam*prob_d
    return mutual_info

def wpmi(clip_feats, target_feats, top_k=28, a=2, lam=0.6, device='cuda', min_prob=1e-7):
    top_k = min(top_k, target_feats.size(0)) 
    with torch.no_grad():
        torch.cuda.empty_cache()
        
        clip_feats = torch.nn.functional.softmax(a*clip_feats, dim=1)

        inds = torch.topk(target_feats, dim=0, k=top_k)[1]
        prob_d_given_e = []

        for orig_id in range(target_feats.shape[1]):
            torch.cuda.empty_cache()
            curr_clip_feats = clip_feats.gather(0, inds[:,orig_id:orig_id+1].expand(-1,clip_feats.shape[1])).to(device)
            curr_p_d_given_e = torch.sum(torch.log(curr_clip_feats+min_prob), dim=0, keepdim=True)
            prob_d_given_e.append(curr_p_d_given_e)

        prob_d_given_e = torch.cat(prob_d_given_e, dim=0)
        #logsumexp trick to avoid underflow
        prob_d = (torch.logsumexp(prob_d_given_e, dim=0, keepdim=True) -
                  torch.log(prob_d_given_e.shape[0]*torch.ones([1]).to(device)))

        mutual_info = prob_d_given_e - lam*prob_d
    return mutual_info

def rank_reorder(clip_feats, target_feats, device="cuda", p=3, top_fraction=0.05, scale_p=0.5):
    """
    top fraction: percentage of mostly highly activating target images to use for eval. Between 0 and 1
    """
    relu = torch.nn.ReLU()
    with torch.no_grad():
        batch = 1500
        errors = []
        top_n = int(target_feats.shape[0]*top_fraction)
        target_feats, inds = torch.topk(target_feats, dim=0, k=top_n)

        for orig_id in range(target_feats.shape[1]):
            clip_indices = clip_feats.gather(0, inds[:, orig_id:orig_id+1].expand([-1,clip_feats.shape[1]])).to(device)
            #calculate the average probability score of the top neurons for each caption
            avg_clip = torch.mean(clip_indices, dim=0, keepdim=True)
            # FIXME
            avg_clip = relu(avg_clip)
            clip_indices = torch.argsort(clip_indices, dim=0)
            curr_errors = []
            target = target_feats[:, orig_id:orig_id+1].to(device)
            sorted_target = torch.flip(target, dims=[0])
            baseline_diff = sorted_target - torch.cat([sorted_target[torch.randperm(len(sorted_target))] for _ in range(5)], dim=1)
            baseline_diff = torch.mean(torch.abs(baseline_diff)**p)
            torch.cuda.empty_cache()
            for i in range(math.ceil(clip_indices.shape[1]/batch)):

                clip_id = (clip_indices[:, i*batch:(i+1)*batch])
                reorg = sorted_target.expand(-1, batch).gather(dim=0, index=clip_id)
                diff = (target-reorg)
                curr_errors.append(torch.nanmean(torch.abs(diff)**p, dim=0, keepdim=True)/baseline_diff)
            errors.append(torch.cat(curr_errors, dim=1)/(avg_clip)**scale_p)

        errors = torch.cat(errors, dim=0)
    return -errors



def max_sim(info_score, concepts, sim):
    return torch.amax(info_score)

def max_minus_sec(info_score, concepts, sim):
    top = torch.topk(info_score, k=2)
    return top.values[0] - top.values[1]

def max_minus_other(info_score, concepts, sim):
    return 2 * torch.amax(info_score) - torch.sum(info_score)

def weighted_diff(info_score, concepts, sim):
    max_val, max_id = torch.topk(info_score, k=1)
    penalty = 0
    for i in range(len(info_score)):
        if i != max_id:
            penalty += info_score[i] * sim[concepts[max_id]][concepts[i]]
    return max_val - penalty