构建自定义损失来评估 PyTorch 在图神经网络上的排名，而无需破坏梯度计算和反向传播

Question

我正在尝试在 PyTorch 中创建一个自定义损失函数，以更好地适应我在 GNN 上的学习任务（我正在使用 PyTorch Geometric 框架）。 只是关于我的学习问题的概述：我有一个 HeteroGraph (https://pytorch-geometric.readthedocs.io/en/latest/notes/heterogeneous.html)，有两种不同类型的节点（我们称它们为 A 和 B ）并且想要预测附加到节点 A 和节点 B 之间的链接的目标（连续值）（知道我也有节点 A 和节点 A 之间的链接，以及节点 B 和节点 B 之间的链接，它们在卷积层中使用，但是我不是在预测他们）。我首先使用均方误差和平均绝对误差，因为它们通常用于评估回归任务的性能，并且最终效果良好。

但是，我想尝试另一种方法，我对值本身不太感兴趣，而是对它们的预测值的顺序感兴趣。然后，我决定使用指标来比较排名：从特定节点，获取所有传出边缘，预测输出，对这些预测进行排名（首先是最高值），对基本事实进行排名并比较这两个排名 - 我所说的排名是指列表按输出/目标值排序时的目标节点数。 为此，我决定使用排名偏向重叠指标（此处定义：https://doi.org/10.1145/1852102.1852106），它具有自定义 python 包 - 但我也可以使用任何评估未标记排名的指标，例如肯德尔相关性。

我想使用这个指标作为执行我的学习任务的损失，但是有一些挑战需要超越（如果我对某些技术方面感到困惑，提前抱歉，我不习惯深入研究深度学习方法）：

• 我想保持 PyTorch 张量梯度的运行，以便 正确执行反向传播
• 我无法直接从一组边计算我的损失，我需要通过以下方式计算它的源节点 源节点（因为只有当我们以节点为中心时排名才相关）
• 我仍然需要根据传出的数量对计算出的 RBO 进行加权 来自模型正在训练的节点的边缘（更容易拥有 在两条边上的排名优于在八条不同边上的排名，并且 RBO 没有反映这一点）
• 由于每个节点都有自己的连接，我无法将所有排名放在一个张量中并使用矩阵运算（排名张量的维度）来计算我的损失 有所不同，某些节点比其他节点连接得更紧密）。
• 我目前 依赖于对包 RBO 的外部调用（尽管它可能是 手动实现）。

我不确定如何在不破坏梯度并使向后传递不可行的情况下实现此自定义损失。如果我理解正确，那么我应该仅通过张量运算直接根据原始输入张量（目标和预测输出）计算损失值？

下面是我的代码的当前状态。返回

evaluate_nodes()

中的两个张量列表并初始化空张量以在前向传递中一个一个地填充它们，这看起来相当难看，但我不确定如何处理前面提到的问题。

RBOL 类损失

class RBOLoss(torch.nn.Module):
def __init__(self, reduction='none') -> None:
    if reduction not in ['mean', 'sum', 'none']:
        raise ValueError('RBO Loss: Reduction `{}` not implemented'.format(reduction))
    self.reduction = reduction
    super(RBOLoss, self).__init__()

"""
We expect to have, for each node, two arrays `prediction` and `target` of the same shape containing respectively the predicted output and the ground truth for each of the edges coming out of this node
"""
def calculate_rbo(self, prediction, target):
    # Rank-Biased Overlap
    gt_ranking = torch.argsort(target, descending=True)
    pred_ranking = torch.argsort(prediction, descending=True)
    # RBO package is not accepting tensors
    return (rbo.RankingSimilarity(gt_ranking.numpy(), pred_ranking.numpy()).rbo(), gt_ranking.size(dim=0))

def forward(self, predictions, targets):
    rbos, weights = torch.empty(len(predictions), dtype=float, requires_grad=True), torch.empty(len(targets), dtype=float, requires_grad=True)
    for idx, (pred, gt) in enumerate(zip(predictions, targets)):
        rbo, weight = self.calculate_rbo(pred, gt)
        rbos[idx] = 1 - rbo # Computing loss, RBO=1 is perfect prediction
        weights[idx] = weight

    loss = torch.mul(rbos, weights)
    if self.reduction == 'mean':
        loss = torch.div(loss.sum(), weights.sum())
    elif self.reduction == 'sum':
        loss = loss.sum()
    return loss

模型中的代码片段

def evaluate_nodes(self, x_dict, edge_index_dict, edge_label_index, targets):
    # Compute output and give ground truths node by node - for ranking purposes
    outputs, ground_truths = [], []
    for node_idx in edge_label_index[0].unique():
        sel = torch.where(edge_label_index[0] == node_idx)
        outputs.append(self.forward(x_dict, edge_index_dict, edge_label_index[:, sel[0]]))
        ground_truths.append(targets[sel[0]])
    return outputs, ground_truths

主管道中的代码片段（model.inspect(data)返回我感兴趣的链接类型的所有边缘）

def train_rbo(data, optimizer):
model.train()
target = model.inspect(data).y.float().view(-1)
predictions, ground_truths = model.evaluate_nodes(data.x_dict, data.edge_index_dict, model.inspect(data).edge_index, target)
criterion = RBOLoss(reduction='mean')
loss = criterion(predictions, ground_truths)
loss.backward()
optimizer.step()
return loss.item()

Answer 1

也许这个 page 可能会有所帮助，您需要通过包含子类

torch.autograd

来扩展

Function

并手动实现

forward()

和

backward()

方法来跟踪向前和向后的梯度通过。