PyTorch Geometric自定义CorrelationLayer参数不更新问题求助

核心问题 #

你的CorrelationLayer权重不更新的原因有两个关键问题：

1. GCNConv未使用邻接矩阵权重：你仅传递了edge_index给GCNConv，但没有传入对应的边权重，导致CorrelationLayer输出的权重矩阵完全未参与GCN计算，梯度无法回传至self.weights。
2. 皮尔逊系数计算脱离计算图：使用scipy的pearsonr会生成非张量的数值，赋值给correlations后会切断梯度传播链路，即使后续使用边权重，梯度也无法传递到可训练参数。

修复步骤 #

1. 让GCNConv使用邻接矩阵的权重 #

修改GCN模型的forward方法，将dense_to_sparse生成的边索引和边权重都传入GCNConv：

class GCN(nn.Module):
    def __init__(self, num_time_series, ts_length, hidden_channels):
        super(GCN, self).__init__()
        self.corr_layer = CorrelationLayer(num_time_series)
        self.graph_conv = GCNConv(ts_length, hidden_channels)

    def forward(self, x):
        adj = self.corr_layer(x)
        # 同时获取边索引和边权重
        edge_index, edge_weight = torch_geometric.utils.dense_to_sparse(adj)
        # 将边权重传入GCNConv
        out = self.graph_conv(x, edge_index, edge_weight)
        return out

2. 实现可微分的皮尔逊系数计算 #

替换scipy的pearsonr，用PyTorch张量操作实现支持自动微分的版本：

def pearsonr_torch(x, y):
    """PyTorch实现的可微分皮尔逊系数计算"""
    mean_x = torch.mean(x)
    mean_y = torch.mean(y)
    xm = x - mean_x
    ym = y - mean_y
    r_num = torch.sum(xm * ym)
    r_den = torch.sqrt(torch.sum(xm ** 2) * torch.sum(ym ** 2))
    # 处理分母为0的情况，限制系数范围在[-1,1]
    r = torch.clamp(r_num / (r_den + 1e-8), -1.0, 1.0)
    return r, None

然后修改CorrelationLayer的forward方法使用该函数：

class CorrelationLayer(nn.Module):
    def __init__(self, num_time_series):
        super().__init__()
        self.num_time_series = num_time_series
        self.weights = nn.Parameter(torch.rand((num_time_series, num_time_series)))
        # 可选：强制权重矩阵对称，符合邻接矩阵的对称性
        self.weights = nn.Parameter((self.weights + self.weights.T) / 2)

    def forward(self, x): 
        correlations = torch.zeros((x.shape[0], x.shape[0]), device=x.device)
        for i in range(x.shape[0]):
            for j in range(i+1, x.shape[0]):
                c, _ = pearsonr_torch(x[i], x[j])
                correlations[i, j] = c
                correlations[j, i] = c
        correlations = correlations * self.weights
        return correlations

3. 可选：向量化优化皮尔逊系数计算 #

双重循环效率较低，改成向量化操作提升性能：

class CorrelationLayer(nn.Module):
    def __init__(self, num_time_series):
        super().__init__()
        self.num_time_series = num_time_series
        self.weights = nn.Parameter(torch.rand((num_time_series, num_time_series)))
        self.weights = nn.Parameter((self.weights + self.weights.T) / 2)

    def forward(self, x): 
        # x shape: [num_nodes, ts_length]
        x_centered = x - x.mean(dim=1, keepdim=True)
        # 计算协方差矩阵
        cov_matrix = x_centered @ x_centered.T
        # 计算标准差矩阵
        std_matrix = torch.sqrt(torch.sum(x_centered**2, dim=1, keepdim=True)) @ torch.sqrt(torch.sum(x_centered**2, dim=1, keepdim=True)).T
        # 计算皮尔逊相关矩阵，加epsilon避免除以0
        corr_matrix = cov_matrix / (std_matrix + 1e-8)
        corr_matrix = torch.clamp(corr_matrix, -1.0, 1.0)
        # 乘以可训练权重
        corr_matrix = corr_matrix * self.weights
        return corr_matrix

额外注意事项 #

• 确保所有张量都在同一设备（CPU/GPU）上运行，避免梯度传播异常。
• 当前设置的学习率0.5过高，建议先尝试0.01~0.1的范围，避免训练不稳定。

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