如何正确计算神经网络中2D特征图的协方差矩阵？

问题根源分析 #

你的初始实现和NumPy结果出现1/2差异的核心原因，在于样本数量的计算错误：
你通过tensor.reshape(bn, nk, 2, -1)把原本w*h个元素按2个一组拆分，得到的样本数实际是(w*h)//2，而非w*h。

NumPy的np.cov默认采用无偏协方差估计，分母是样本数-1；而你初始代码里误用了h*w -1作为分母，这个数值恰好是正确样本数的2倍左右，最终导致结果是NumPy的1/2。

修正后的实现 #

你已经找到正确的分母计算方式，这里给出优化后的代码（简化维度操作，减少冗余计算）：

import torch

def get_covariance(tensor):
    bn, nk, w, h = tensor.shape
    # 按2个元素一组reshape，得到(2, N)结构，N为有效样本数
    tensor_reshape = tensor.reshape(bn, nk, 2, -1)
    x, y = tensor_reshape[:, :, 0, :], tensor_reshape[:, :, 1, :]
    
    # 计算均值，用keepdim=True避免维度丢失
    mean_x = torch.mean(x, dim=2, keepdim=True)
    mean_y = torch.mean(y, dim=2, keepdim=True)
    
    # 计算无偏协方差，分母为样本数-1
    sample_num = (w * h) // 2
    xx = torch.sum((x - mean_x) ** 2, dim=2, keepdim=True) / (sample_num - 1)
    xy = torch.sum((x - mean_x) * (y - mean_y), dim=2, keepdim=True) / (sample_num - 1)
    yy = torch.sum((y - mean_y) ** 2, dim=2, keepdim=True) / (sample_num - 1)
    
    # 拼接成2x2协方差矩阵，yx与xy对称，无需重复计算
    cov = torch.cat((xx, xy, xy, yy), dim=2).reshape(bn, nk, 2, 2)
    return cov

验证测试 #

用你提供的测试代码验证，结果会完全匹配：

import numpy as np

a = torch.randn(1, 1, 64, 64)
a_numpy = a.reshape(1, 1, 2, -1).numpy()
torch_cov = get_covariance(a)
numpy_cov = np.cov(a_numpy[0][0])

print("PyTorch协方差结果：")
print(torch_cov)
print("\nNumPy协方差结果：")
print(numpy_cov)

额外优化提示 #

• 用keepdim=True替代unsqueeze(-1)，代码更简洁且不易出错；
• 利用协方差矩阵的对称性，复用xy结果，减少不必要的计算。

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