hackathonf23-Creativity_Underflow/tests/loss1.py

import torch
import numpy as np
import torch.nn.functional as F
import tqdm 


torch.manual_seed(4)

# very inefficient but just to check values
b = 64
k = 64
l = 20
n = 196

theta = -0.5
tau = 0.1

x = torch.randn(b, k, n, device='cuda')
y = torch.randn(b, l, n, device='cuda')

x = F.normalize(x, dim=-1)
y = F.normalize(y, dim=-1)

lp = 0
lq = 0
loss = 0

def exp(x):
    return torch.exp(x)

def p_k(x, y):
    p = 0
    count_p = 0
    for k_i in range(k):
        p_i = 0
        for j in range(l):
            cos_sim = F.cosine_similarity(x[k_i], y[j], dim=-1)
            p_i = max(p_i, cos_sim)
        if p_i > theta:
            p += p_i
            count_p += 1
    return p / (count_p + 1e-5)

for b_i in tqdm.tqdm(range(b)):
    p_ii = exp(p_k(x[b_i], y[b_i]) / tau)
    p_ij = 0
    for b_j in range(b):
        p_j = exp(p_k(x[b_i], y[b_j]) / tau)
        p_ij += p_j
    print(torch.log(p_ii / p_ij))
    lp += torch.log(p_ii / p_ij)
lp /= b
lp *= -1
print(f"lossq: {lp}")
quit()
def q_k(y, x):
    q = 0
    count_q = 0
    for l_j in range(l):
        q_i = 0
        for i in range(k):
            cos_sim = F.cosine_similarity(x[i].unsqueeze(0), y[l_j].unsqueeze(0), dim=-1)
            # print(cos_sim)
            q_i = max(q_i, cos_sim)
        if q_i > theta:
            q += q_i
            count_q += 1
    return q / (count_q + 1e-5)

q_ijs = []
for b_i in range(b):
    q_ii = exp(q_k(y[b_i], x[b_i]) / tau)
    q_ij = 0
    for b_j in range(b):
        q_j = exp(q_k(y[b_i], x[b_j]) / tau)
        q_ij += q_j
        
    q_ijs.append(q_ii)
    lq += q_ii / q_ij

lq /= b
lq *= -1

loss = 0.5 * (lp + lq)

print(lp)
print(lq)
print(loss)