metric_learning_NMI.py

# evaluate the clustering performance
from sklearn.cluster import KMeans
from sklearn.metrics.cluster import normalized_mutual_info_score
import numpy as np

def evaluation(X, Y):
    classN = np.max(Y)+1
    kmeans = KMeans(n_clusters=classN).fit(X)
    nmi = normalized_mutual_info_score(Y, kmeans.labels_, average_method='arithmetic')
    return nmi

output = model(input)
testdata = torch.cat((testdata, output.cpu()), 0)
testlabel = torch.cat((testlabel, target))
nmi = evaluation(testdata.numpy(), testlabel.numpy())