如何设置numpy tensordot的axes参数以计算平均角度误差

Hey there! Let's get your vectorized angular error calculation working properly—no more slow loops needed.

First, let's break down why your np.tensordot approach was throwing that shape mismatch error. Your loop version computes the dot product per (h,w) pair of 2D vectors, meaning for each position (i,j), you're taking the dot of the two 2-element vectors. But when you used axes=2, numpy tried to contract the last two dimensions of both arrays, which doesn't align with your (h,w,2) shape.

Correct np.tensordot Implementation #

To get the per-element dot product, you need to tell tensordot to only contract the last dimension (axis 2) of both arrays. Here's how to adjust the axes parameter:

import numpy as np

def average_angular_error_vec(estimated_oc : np.array, target_oc : np.array):
    estimated_oc = np.float64(estimated_oc)
    target_oc = np.float64(target_oc)
    
    # Compute norms and normalize vectors (use keepdims to avoid reshaping)
    norm1 = np.linalg.norm(estimated_oc, axis=2, keepdims=True)
    norm2 = np.linalg.norm(target_oc, axis=2, keepdims=True)
    unit_vector1 = estimated_oc / norm1
    unit_vector2 = target_oc / norm2
    
    # Correct tensordot axes: contract the last axis (axis 2) of both arrays
    dot_product = np.tensordot(unit_vector1, unit_vector2, axes=((2), (2)))
    # Handle numerical instability (clip to [-1,1] to avoid arccos NaNs)
    dot_product = np.clip(dot_product, -1.0, 1.0)
    
    angle = np.arccos(dot_product)
    return np.mean(angle)

Even Simpler Alternative: Element-wise Multiply + Sum #

Honestly, tensordot is overkill here. A more intuitive way to compute the per-element dot product is to multiply the two arrays element-wise, then sum along the last axis. This does exactly what your loop does, but in vectorized form:

def average_angular_error_vec_simple(estimated_oc : np.array, target_oc : np.array):
    estimated_oc = np.float64(estimated_oc)
    target_oc = np.float64(target_oc)
    
    norm1 = np.linalg.norm(estimated_oc, axis=2, keepdims=True)
    norm2 = np.linalg.norm(target_oc, axis=2, keepdims=True)
    unit_vector1 = estimated_oc / norm1
    unit_vector2 = target_oc / norm2
    
    # Element-wise multiply, then sum over the last axis (2D vector dot product)
    dot_product = np.sum(unit_vector1 * unit_vector2, axis=2)
    dot_product = np.clip(dot_product, -1.0, 1.0)
    
    angle = np.arccos(dot_product)
    return np.mean(angle)

Key Notes #

1. keepdims=True: This removes the need for manual reshaping with [..., np.newaxis]—it keeps the norm array shape as (h,w,1), so broadcasting works seamlessly when dividing the original (h,w,2) arrays.
2. Numerical Stability: The np.clip step is non-negotiable! Floating-point precision errors can make dot_product slightly exceed 1.0 or drop below -1.0, which would cause np.arccos to return nan. Clipping fixes this edge case.
3. Performance: Both vectorized functions will produce the exact same result as your loop version, but run orders of magnitude faster for large arrays—perfect for your (436,1024) shape.

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