reallsp / saf Goto Github PK

def compute_topk(
    query, gallery, target_query, target_gallery, k=[1, 10], reverse=False
):
    result = []
    query0 = query[:,0] / query[:,0].norm(dim=1, keepdim=True)
    gallery0 = gallery[:,0] / gallery[:,0].norm(dim=1, keepdim=True)
    sim_cosine = torch.matmul(query0, gallery0.t())
    score = torch.zeros((sim_cosine.shape))
    for i in range(query.size(1)):
        query0 = query[:,i] / query[:,i].norm(dim=1, keepdim=True)
        gallery0 = gallery[:,i] / gallery[:,i].norm(dim=1, keepdim=True)
        sim_cosine = torch.matmul(query0, gallery0.t())
        score+=sim_cosine
    result.extend(topk(score, target_gallery, target_query, k, dim=1))
    if reverse:
        result.extend(topk(score, target_query, target_gallery, k, dim=0))
    return result, score



def topk(sim, target_gallery, target_query, k=[1, 10], dim=1):
    result = []
    maxk = max(k)
    size_total = len(target_query)
    _, pred_index = sim.topk(maxk, dim, True, True)  # 得到相似度最大的前k个
    pred_labels = target_gallery[pred_index]
    if dim == 1:
        pred_labels = pred_labels.t()
    correct = pred_labels.eq(target_query.view(1, -1).expand_as(pred_labels))

    for topk in k:
        correct_k = torch.sum(correct[:topk], dim=0)
        correct_k = torch.sum(correct_k > 0).float()
        result.append(correct_k * 100 / size_total)
    return result