pytorch NT-Xent代码实现

pytorch实现NT-Xent算法的代码如下：

import torch
import torch.nn.functional as F
from torch import nn

class NTXentLoss(nn.Module):
    def __init__(self, batch_size, temperature):
        super(NTXentLoss, self).__init__()
        self.batch_size = batch_size
        self.temperature = temperature
        self.mask = self.mask_correlated_samples(batch_size)

    def mask_correlated_samples(self, batch_size):
        mask = torch.ones((2 * batch_size, 2 * batch_size), dtype=bool)
        mask = mask.fill_diagonal_(0)
        for i in range(batch_size):
            mask[i, batch_size + i] = 0
            mask[batch_size + i, i] = 0
        return mask

    def forward(self, zis, zjs):
        representations = torch.cat([zis, zjs], dim=0)
        similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
        similarity_matrix = similarity_matrix / self.temperature
        labels = torch.arange(2 * self.batch_size).to(similarity_matrix.device)
        loss = F.cross_entropy(similarity_matrix, labels, reduction='sum')
        loss /= (2 * self.batch_size)
        return loss

# 使用示例
criterion = NTXentLoss(batch_size=64, temperature=0.5)
embedding_size = 128
zis = torch.randn(64, embedding_size)
zjs = torch.randn(64, embedding_size)
loss = criterion(zis, zjs)
print(loss)