evaluate.py

'''
Created on Sep 6, 2021
PyTorch Implementation of Multi-Layer Perceptron recommender model in:
He Xiangnan et al. Neural Collaborative Filtering. In WWW 2017.

@author: Yitong Wang (yitongwang.snowflake@gmail.com)
'''

import math
import heapq#for retrieval top_K
import multiprocessing as mp
import torch
import numpy as np
from time import time

#Global variables that are shared across processes

_model=None
_testRatings=None
_testNegatives=None
_K=None
_device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

def evaluate_model(model,testRatings,testNegatives,K,num_thread):
    """

    Evaluate the performance (Hit_Ratio, NDCG) of top_K recommendation
    Return: score of each test rating.
    """
    global _model
    global _testRatings
    global _testNegatives
    global _K
    _model=model.to(_device)
    _testRatings=testRatings
    _testNegatives=testNegatives
    _K=K

    hits,ndcgs=[],[]
    if num_thread>1:
        pool=mp.Pool(processes=num_thread)
        res=pool.map(eval_one_rating,range(len(_testRatings)))
        pool.close()
        pool.join()
        hits=[r[0] for r in res]
        ndcgs=[r[1] for r in res]
        return (hits,ndcgs)
    #single thread
    for idx in range(len(_testRatings)):
        (hr,ndcg)=eval_one_rating(idx)
        hits.append(hr)
        ndcgs.append(ndcg)
    return (hits,ndcgs)

def eval_one_rating(idx):
    rating=_testRatings[idx]
    items=_testNegatives[idx]
    u=rating[0]
    gtItem=rating[1]
    items.append(gtItem)
    #Get prediction scores
    map_item_score={}
    users=np.full(len(items),u,dtype='int32')#numpy.full(shape, fill_value, dtype=None, order=‘C’)返回一个形为shape的数组
    users,items=torch.tensor(users),torch.tensor(items)
    users,items=users.to(_device),items.to(_device)
    predictions=_model(users,items)
    items=items.to('cpu')
    predictions=predictions.to('cpu')
    for i in range(len(items)):
        item=items[i]
        map_item_score[item]=predictions[i]
    items=list(items)
    items.pop()

    ranklist=heapq.nlargest(_K,map_item_score,key=map_item_score.get)
    hr=getHitRatio(ranklist,gtItem)
    ndcg=getNDCG(ranklist,gtItem)
    return (hr,ndcg)

def getHitRatio(ranklist,gtItem):
    for item in ranklist:
        if item==gtItem:
            return 1
    return 0

def getNDCG(ranklist,gtItem):
    for i in range(len(ranklist)):
        item=ranklist[i]
        if item==gtItem:
            return math.log(2)/math.log(i+2)
    return 0