矩阵的余弦距离矩阵，GPU加速

Question

我真诚地请求您的帮助。我已经被这个问题困扰了2天多了。简单来说，当我想计算向量的距离时，我只需简单地使用：

from sklearn.metrics.pairwise import cosine_distances
distances = cosine_distances(vector_matrix)

其中

vector_matrix

的形状为 [4,768]。我想计算 4 个向量的成对余弦距离，每个向量有 768 个维度。然后我就可以得到距离矩阵，一个 4*4 的对称矩阵。

上周我使用 sciBERT 嵌入了一些文本，然后每个文本被编码成一个 512*768 的矩阵。现在我的

embeded_text

是 [4, 512, 768]。

cosine_distances

只接受向量而不接受矩阵。有没有一种简单的方法来计算矩阵之间的余弦距离矩阵（就像

cosine_distances

）？我可以一一计算矩阵余弦距离，但我的真实数据是[10663, 512, 768]，GPU加速是必须的。

Answer 1

根据

cosine_distances

的sklearn文档，它们的输入参数是 第一个特征的

(n_samples_X, n_features)

的形状和第二个特征的

(n_samples_Y, n_features)

的形状，在这种情况下，

n_samples

是你的批量大小（假设它是你的数据，4 是你的

n_samples

）。

然后，我们需要考虑的是使用 sciBERT 嵌入文本的

(4, 512, 768)

形状。首先，正如我所说，我们已经有 4 个，即

n_samples

。其余的是

n_features

，它需要一个特征向量作为余弦距离计算的输入，但我们有一个包含 (512, 768) 的矩阵。这可以通过将特征矩阵重塑为特征向量来解决，而不是通过

reshape

函数。

# initialize a matrix with batch size 4    
matrix = np.random.rand(4, 512, 768)

# reshape into the shape (n_samples_X, n_features) required for cosine distances function, where n_features will equal to 512*768 and represented as a feature vector
reshaped_matrix = matrix.reshape(4, -1) # shape: (4, 393216)

print(cosine_distances(reshaped_matrix)) # output shape: (4, 4)

这样，我们仍然可以使用

cosine_distances

函数，就像我们使用之前的批量向量来生成 cosine_distances 值一样。