numpy 实现nt-xent

下面是使用numpy实现NT-Xent Loss的示例代码：

import numpy as np

def nt_xent_loss(embeddings, temperature=0.5):
    # embeddings: shape (2N, D), 其中N是样本数，D是特征向量的维度
    # temperature: 温度参数，用于调整相似度的尺度
    N = embeddings.shape[0] // 2
    # 计算余弦相似度矩阵
    sims = np.dot(embeddings, embeddings.T)
    # 将对角线置为0，避免将相同的样本作为负样本
    sims = sims - np.eye(2*N)
    # 计算每个样本的相似度排名，并将排名转换为概率分布
    pos_sims = np.diag(sims[:N, N:])
    pos_probs = pos_sims / temperature
    neg_sims = np.exp(sims[:N, :N] / temperature)
    neg_probs = neg_sims / np.sum(neg_sims, axis=1, keepdims=True)
    # 计算NT-Xent Loss
    loss = -np.mean(np.log(pos_probs / (pos_probs + np.sum(neg_probs, axis=1))))
    return loss

其中，输入的embeddings是一个形状为(2N, D)的numpy数组，其中N是样本数，D是特征向量的维度。函数首先计算余弦相似度矩阵，然后将对角线置为0，避免将相同的样本作为负样本。接下来，计算每个样本的相似度排名，并将排名转换为概率分布。最后，计算NT-Xent Loss并返回。参数temperature用于调节相似度的尺度，通常取0.5到1之间的值。