使用tf.map_fn调用Scipy距离函数时出现张量形状不兼容错误

核心原因 #

1. tf.py_function形状推断错误：训练时TensorFlow处于图模式，若未明确指定output_signature，TensorFlow无法正确识别Scipy距离函数返回的标量形状，会误判为与输入同维度的张量（即[5000]），最终和期望的[batch_size]损失张量形状不兼容。小批次测试用的是eager模式，形状推断宽松，因此未暴露问题。
2. Scipy函数调用逻辑有误：你的输入是[num_classes]的概率向量（离散分布的概率质量），但Scipy的wasserstein_distance默认处理样本集而非离散分布。直接传入概率向量时，函数会将每个元素视为独立样本点，虽返回标量，但图模式下形状推断易出错。
3. 自动微分兼容性问题：tf.py_function本身不支持自动梯度回传，训练时梯度计算过程中，形状不匹配的问题会被放大，触发崩溃。

解决步骤 #

1. 修正Scipy距离函数的调用方式（针对离散概率分布） #

针对每个样本的概率向量p_i和q_i，需明确传入类别索引作为分布的取值点，确保计算的是离散分布间的距离：

from scipy.stats import wasserstein_distance
from scipy.spatial.distance import energy
import numpy as np

num_classes = 5000
# 定义类别索引，对应每个概率的位置
class_indices = np.arange(num_classes)

def compute_wasserstein(p_np, q_np):
    # 计算离散分布的Wasserstein距离
    return wasserstein_distance(class_indices, class_indices, 
                                u_weights=p_np, v_weights=q_np)

def compute_energy_distance(p_np, q_np):
    # 计算离散分布的能量距离
    return energy(class_indices, class_indices, 
                  u_weights=p_np, v_weights=q_np)

2. 明确tf.py_function的output_signature #

封装函数时必须指定返回张量为标量，让TensorFlow在图模式下正确推断形状：

import tensorflow as tf

def tf_custom_loss(p, q, distance_func):
    def py_wrap(p_i, q_i):
        p_np = p_i.numpy()
        q_np = q_i.numpy()
        return distance_func(p_np, q_np)
    
    # 遍历batch维度，每个元素是[num_classes]的向量
    loss_per_sample = tf.map_fn(
        lambda x: tf.py_function(
            func=py_wrap,
            inp=[x[0], x[1]],
            Tout=tf.float32,
            # 明确指定返回标量张量
            output_signature=tf.TensorSpec(shape=(), dtype=tf.float32)
        ),
        elems=(p, q),
        dtype=tf.float32
    )
    # 返回批次平均损失
    return tf.reduce_mean(loss_per_sample)

# 使用示例：Wasserstein损失
wasserstein_loss = tf_custom_loss(p, q, compute_wasserstein)

3. 支持自动梯度回传（训练必备） #

tf.py_function不支持自动微分，需用tf.custom_gradient封装梯度计算逻辑（以下为数值梯度示例，也可推导解析梯度）：

def wasserstein_with_grad(p, q):
    @tf.custom_gradient
    def _compute(p, q):
        p_np = p.numpy()
        q_np = q.numpy()
        loss = compute_wasserstein(p_np, q_np)
        
        def grad(dy):
            eps = 1e-6
            # 计算对p的梯度
            grad_p = np.zeros_like(p_np)
            for i in range(num_classes):
                p_eps = p_np.copy()
                p_eps[i] += eps
                loss_eps = compute_wasserstein(p_eps, q_np)
                grad_p[i] = (loss_eps - loss) / eps
            # 计算对q的梯度
            grad_q = np.zeros_like(q_np)
            for i in range(num_classes):
                q_eps = q_np.copy()
                q_eps[i] += eps
                loss_eps = compute_wasserstein(p_np, q_eps)
                grad_q[i] = (loss_eps - loss) / eps
            # 转换为TensorFlow张量并乘以梯度流
            return dy * tf.convert_to_tensor(grad_p, tf.float32), dy * tf.convert_to_tensor(grad_q, tf.float32)
        
        return tf.convert_to_tensor(loss, tf.float32), grad
    
    return _compute(p, q)

# 封装到map_fn中
def tf_wasserstein_loss_with_grad(p, q):
    loss_per_sample = tf.map_fn(
        lambda x: wasserstein_with_grad(x[0], x[1]),
        elems=(p, q),
        dtype=tf.float32
    )
    return tf.reduce_mean(loss_per_sample)

4. 替代方案：使用TensorFlow原生实现（推荐） #

避免使用Scipy函数，改用TensorFlow Probability库的原生实现，完全兼容图模式和自动微分：

import tensorflow_probability as tfp

def tf_wasserstein_loss_tfp(p, q):
    # 将概率向量转换为Categorical分布
    dist_p = tfp.distributions.Categorical(probs=p)
    dist_q = tfp.distributions.Categorical(probs=q)
    # 计算Wasserstein距离（order=1）
    loss_per_sample = tfp.distributions.wasserstein_distance(dist_p, dist_q, order=1)
    return tf.reduce_mean(loss_per_sample)

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