基于nDCG实现PyTorch损失函数时遇RuntimeError问题求助

问题背景 #

基于nDCG实现的PyTorch损失函数代码如下：

import torch

class NDCGLoss(torch.nn.Module):

    def __init__(self, relevance_map):
        super(NDCGLoss, self).__init__()
        self.relevance_map = relevance_map

    def get_relevance(self, queries_ids, docs_ids):
        relevance = torch.empty((queries_ids.shape[0], docs_ids.shape[0]), device=queries_ids.device)
        for i, query_idx in enumerate(queries_ids.tolist()):
            for j, doc_idx in enumerate(docs_ids.tolist()):
                if doc_idx in self.relevance_map[query_idx]:
                    relevance[i, j] = 1.0
                else:
                    relevance[i, j] = 0.0
        return relevance

    def _get_dcg(self, scores, relevances):
        discount = 1.0 / (torch.log2(torch.arange(relevances.shape[-1], device=relevances.device) + 2.0))
        ranking = scores.argsort(descending=True)
        ranked = torch.gather(relevances, dim=-1, index=ranking)
        return torch.sum(discount * ranked, dim=-1)

    def forward(self, query_idx, query_rpr, doc_idx, doc_rpr):

        scores = torch.einsum("ab,cb->ac", query_rpr, doc_rpr)  # inner product
        relevance = self.get_relevance(query_idx, doc_idx)  # ground-truth
        dcg = self._get_dcg(scores, relevance)
        idcg = self._get_dcg(relevance, relevance)
        idcg = torch.where(idcg == 0, 1.0, idcg)
        return -torch.log(
            torch.mean(
                torch.div(dcg, idcg)
            ) + 1e-11
        )

运行时遇到两类RuntimeError：

1. RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.
2. 执行output.backward()时触发：RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

测试用例代码：

query_idx = torch.tensor([1,2,3])
query_rpr = torch.rand((3,5), requires_grad=True)

doc_idx = torch.tensor([4,5,6])
doc_rpr = torch.rand((3,5), requires_grad=True )


relevance_map = {
    1: [4,8],
    2: [5,6],
    3: [7,9]
}

loss = NDCGLoss(relevance_map)
output = loss(query_idx, query_rpr, doc_idx, doc_rpr)
output.backward()

错误原因分析 #

• in-place操作错误：get_relevance中用torch.empty创建的relevance张量会继承输入张量的requires_grad属性，后续循环逐元素赋值属于in-place操作，而叶子节点的视图不允许此类修改。
• 梯度传播中断：相关性标签张量未明确设置requires_grad=False，导致后续计算中梯度传播路径出现无梯度的张量，最终损失无法反向传播。

修复方案及完整代码 #

关键修复点 #

1. 相关性张量明确设置requires_grad=False，用向量化操作替代循环赋值，避免in-place操作。
2. 用torch.clamp替代torch.where处理IDCG为0的边界情况，保持计算图连续性。
3. 提前将相关性映射转换为张量，提升计算效率。

修复后的代码：

import torch

class NDCGLoss(torch.nn.Module):

    def __init__(self, relevance_map):
        super(NDCGLoss, self).__init__()
        self.relevance_map = relevance_map
        # 提前转换为张量映射，减少循环内类型转换开销
        self.query_doc_tensor_map = {k: torch.tensor(v, dtype=torch.long) for k, v in relevance_map.items()}

    def get_relevance(self, queries_ids, docs_ids):
        # 创建不需要梯度的相关性矩阵，避免in-place操作问题
        relevance = torch.zeros((queries_ids.shape[0], docs_ids.shape[0]), 
                               device=queries_ids.device, 
                               dtype=torch.float32,
                               requires_grad=False)
        
        # 向量化匹配替代循环赋值
        for i, q_id in enumerate(queries_ids.tolist()):
            matched_mask = torch.isin(docs_ids, self.query_doc_tensor_map[q_id])
            relevance[i] = matched_mask.float()
        
        return relevance

    def _get_dcg(self, scores, relevances):
        # 修正折扣项的索引计算，从1开始对应排名位置
        discount = 1.0 / torch.log2(torch.arange(1, relevances.shape[-1]+1, device=relevances.device) + 1.0)
        ranking = scores.argsort(descending=True)
        ranked_relevances = torch.gather(relevances, dim=-1, index=ranking)
        return torch.sum(discount * ranked_relevances, dim=-1)

    def forward(self, query_idx, query_rpr, doc_idx, doc_rpr):
        scores = torch.einsum("ab,cb->ac", query_rpr, doc_rpr)
        relevance = self.get_relevance(query_idx, doc_idx)
        
        dcg = self._get_dcg(scores, relevance)
        idcg = self._get_dcg(relevance, relevance)
        # 用clamp处理IDCG为0的情况，避免除以0同时保持计算图连续
        idcg = torch.clamp(idcg, min=1e-11)
        
        ndcg = dcg / idcg
        return -torch.log(torch.mean(ndcg) + 1e-11)

测试验证 #

运行原测试用例，可正常计算梯度：

query_idx = torch.tensor([1,2,3])
query_rpr = torch.rand((3,5), requires_grad=True)

doc_idx = torch.tensor([4,5,6])
doc_rpr = torch.rand((3,5), requires_grad=True )

relevance_map = {
    1: [4,8],
    2: [5,6],
    3: [7,9]
}

loss = NDCGLoss(relevance_map)
output = loss(query_idx, query_rpr, doc_idx, doc_rpr)
output.backward()

# 验证梯度是否正常生成
print(query_rpr.grad is not None)  # 输出True
print(doc_rpr.grad is not None)    # 输出True

内容的提问来源于stack exchange，提问作者Celso França