python - 特定轴上 np.corrcoef 的 Tensorflow 等效项

Question

我正在尝试按列关联两个矩阵。即，将第一个矩阵的第一列与第二个矩阵的第一列相关联，依此类推。在 numpy 中我这样做:

np.corrcoef(x, y, axis=0)

效果很好。该命令的 Tensorflow 等效项是什么？

我尝试使用streaming_pearson_correlation 1但这将所有列关联在一起，而不是为每列提供结果。

作为最后的手段，我正在考虑将张量拆分为单独的列张量，但我猜测这会产生性能成本。

我知道我可以将 numpy 包装在 py_func 中，但它不会在 GPU 上运行。

预先感谢您的帮助。

Answer 1

numpy 文档页面 corrcoef给出了 corcoef 和协方差矩阵之间的联系.因此，自然的事情是首先用 numpy 中的 matmuls 重写它:

fsize=1
dsize=3
x=np.random.random((fsize,dsize))
y=np.random.random((fsize,dsize))
xy=np.concatenate([x,y], axis=0)
(np.corrcoef(xy) == np.corrcoef(x,y)).all()
mean = np.mean(xy, axis=1, keepdims=True)
cov = ((xy-mean) @ (xy-mean).T)/(dsize-1)
cov2 = np.diag(1/sqrt(np.diag(cov)))
np.testing.assert_allclose(cov2@cov@cov2, np.corrcoef(x, y))

现在转换为 TensorFlow，并检查结果是否相同

def t(x): return tf.transpose(x)
sess = tf.InteractiveSession()

x_t = tf.constant(x)
y_t = tf.constant(y)
xy_t = tf.concat([x, y], axis=0)
mean_t = tf.reduce_mean(xy_t, axis=1, keep_dims=True)
cov_t = ((xy_t-mean_t) @ t(xy_t-mean_t))/(dsize-1)
cov2_t = tf.diag(1/tf.sqrt(tf.diag_part(cov_t)))
cor = cov2_t @ cov_t @ cov2_t

np.testing.assert_allclose(np.corrcoef(x, y), cor.eval())

构成 x 和 y 的变量之间的相关性位于该矩阵的非对角 block 中。