如何高效计算3D与2D NumPy数组间的余弦相似度？

Great question! That loop works, but it can get pretty slow when m is large—Python loops add overhead, and you're calling cosine_similarity m times instead of doing the computation in one go. Let's fix this with a fully vectorized NumPy solution that leverages optimized linear algebra operations, which will be way faster.

How It Works #

Cosine similarity between two vectors u and v is calculated as:

cos_sim(u, v) = (u · v) / (||u|| * ||v||)

For your arrays:

• A is shape (m, n, 300): each (n, 300) sub-matrix has n vectors of length 300
• B is shape (p, 300): p vectors of length 300

We can compute all pairwise cosine similarities across all m sub-matrices in one vectorized step:

1. Compute L2 norms for all vectors in A and B (we need these for the denominator)
2. Calculate dot products between every vector in A and every vector in B
3. Normalize the dot products by the product of the norms to get cosine similarities

Code Implementation #

import numpy as np

# Assume A is (m, n, 300) and B is (p, 300)
m, n, dim = A.shape
p = B.shape[0]

# Step 1: Compute L2 norms (keepdims preserves shape for broadcasting)
A_norm = np.linalg.norm(A, axis=2, keepdims=True)  # Shape: (m, n, 1)
B_norm = np.linalg.norm(B, axis=1, keepdims=True)  # Shape: (p, 1)

# Step 2: Compute dot products across all vectors
# A @ B.T gives (m, n, p) since B.T is (300, p)
dot_product = A @ B.T

# Step 3: Calculate cosine similarity, handle division by zero
denominator = A_norm @ B_norm.T  # Shape: (m, n, p)
# Replace 0s in denominator to avoid division errors (adjust if needed for your use case)
denominator[denominator == 0] = 1
cos_sim = dot_product / denominator

# If you need the exact same format as your original loop result (list of (n,p) arrays):
result = [cos_sim[i] for i in range(m)]

Why This Is Better #

• Speed: Vectorized operations use NumPy's optimized C-backed linear algebra routines (like BLAS/LAPACK) instead of slow Python loops. For large m, this can be 10-100x faster.
• Memory Efficiency: It avoids creating intermediate lists and reuses array memory more effectively.
• Consistency: You get the exact same values as the sklearn loop (we're replicating the same cosine similarity calculation under the hood).

Notes #

• If you have vectors with zero norm (all zeros), the code replaces the denominator with 1 to avoid division by zero. You might want to handle these cases differently (e.g., set similarity to 0) depending on your data.
• This approach uses more memory than the loop (since it computes all similarities at once), but for most practical cases, it's worth the tradeoff for speed.

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