如何在Pytorch中计算2D和3D张量之间的欧氏距离

Question

鉴于：

• 张量 A 具有形状

  (batch_size, dim)

• 张量 B 具有形状

  (batch_size, N, dim)

我想计算 A 中的每一行与 B 中形状为

(N, dim)

的相应行之间的欧氏距离

预期结果已成型

(batch_size, N)

Answer 1

你可以试试这个；它帮助了我，也可能解决你的问题。

import numpy as np
def euclidean_distance(A, B):

A_expanded = np.expand_dims(A, axis=1)

squared_diff = np.square(A_expanded - B)

squared_distance = np.sum(squared_diff, axis=-1)

euclidean_distance = np.sqrt(squared_distance)

return euclidean_distance

batch_size = 3
N = 4
dim = 2

# Generate sample data
A = np.random.rand(batch_size, dim)
B = np.random.rand(batch_size, N, dim)

# Calculate Euclidean distance
distances = euclidean_distance(A, B)
print(distances)