constrastive_loss.py

def loss(a,b):
    a_norm = torch.norm(a,dim=1).reshape(-1,1)
    a_cap = torch.div(a,a_norm)
    b_norm = torch.norm(b,dim=1).reshape(-1,1)
    b_cap = torch.div(b,b_norm)
    a_cap_b_cap = torch.cat([a_cap,b_cap],dim=0)
    a_cap_b_cap_transpose = torch.t(a_cap_b_cap)
    b_cap_a_cap = torch.cat([b_cap,a_cap],dim=0)
    sim = torch.mm(a_cap_b_cap,a_cap_b_cap_transpose)
    sim_by_tau = torch.div(sim,tau)
    exp_sim_by_tau = torch.exp(sim_by_tau)
    sum_of_rows = torch.sum(exp_sim_by_tau, dim=1)
    exp_sim_by_tau_diag = torch.diag(exp_sim_by_tau)
    numerators = torch.exp(torch.div(torch.nn.CosineSimilarity()(a_cap_b_cap,b_cap_a_cap),tau))
    denominators = sum_of_rows - exp_sim_by_tau_diag
    num_by_den = torch.div(numerators,denominators)
    neglog_num_by_den = -torch.log(num_by_den)
    return torch.mean(neglog_num_by_den)