基于PyTorch与PyG的异规模图相似度矩阵批量计算及损失求解

核心思路 #

直接利用PyG批量数据的batch分割能力，避免填充操作，对每个子图独立计算损失后累加，完全适配PyG的连续张量批量表示，同时最大化GPU并行效率。

实现步骤与代码 #

1. 分割批量特征张量：通过batch张量统计每个子图的节点数，用torch.split拆分出单个图的特征矩阵。
2. 计算相似度矩阵：对每个子图特征矩阵计算X_i @ X_i.T作为交叉熵的输入logits。
3. 构造目标与计算损失：每个子图的目标是节点自身索引（对应单位矩阵的正样本），直接用F.cross_entropy计算单图损失后累加平均。

import torch
import torch.nn.functional as F
from torch_geometric.data import Batch

def batch_intra_graph_loss(x, batch):
    # 统计每个图的节点数量
    node_counts = torch.bincount(batch)
    # 拆分批量特征为各子图的特征矩阵
    subgraph_features = torch.split(x, node_counts.tolist())
    
    total_loss = 0.0
    for x_i in subgraph_features:
        n_nodes = x_i.size(0)
        # 计算节点相似度矩阵（作为交叉熵的logits）
        sim_matrix = x_i @ x_i.T
        # 目标：每个节点的正样本是自身，对应索引序列
        target = torch.arange(n_nodes, device=x.device)
        # 计算单图交叉熵损失
        loss = F.cross_entropy(sim_matrix, target)
        total_loss += loss
    
    # 返回批量平均损失
    return total_loss / len(subgraph_features)

# 示例使用（GPU环境）
# 构造批量数据：3个图，节点数分别为2、3、4，特征维度16
batch = torch.tensor([0,0,1,1,1,2,2,2,2], device='cuda')
x = torch.randn(9, 16, device='cuda')
loss = batch_intra_graph_loss(x, batch)
print(loss.item())

优化说明 #

• 无填充冗余：完全基于PyG原生的连续批量张量，避免填充带来的内存浪费和无效计算。
• GPU并行高效：每个子图的计算独立，PyTorch会自动调度GPU异步执行，循环开销可忽略，批量100的场景下性能远优于单图循环。
• 数值稳定性：若相似度矩阵数值波动大，可先对特征做L2归一化：x_i = F.normalize(x_i, dim=1)，避免交叉熵计算溢出。

内容的提问来源于stack exchange，提问作者Adrien Lagesse