如何在TensorFlow中实现张量的分段中位数计算？

TensorFlow没有内置的tf.math.segment_median函数，但可以通过组合现有操作实现，核心步骤如下：

• 按分段ID拆分原始张量
• 对每个分段内的元素沿样本维度排序
• 根据分段长度计算中位数位置，提取对应值（奇数长度取中间元素，偶数长度取中间两数的平均值）
• 拼接各分段的中位数得到最终结果

代码实现 #

import tensorflow as tf

x = tf.constant([[0.1, 0.2, 0.4, 0.5],
                 [0.1, 0.8, 0.2, 0.6],
                 [0.1, 0.2, 0.4, 0.5],
                 [0.3, 0.1, 0.2, 0.9],
                 [0.1, 0.1, 0.6, 0.5]])
segment_ids = tf.constant([0, 1, 1, 1, 2])

def segment_median(x, segment_ids):
    # 按分段ID拆分张量
    partitions = tf.dynamic_partition(x, segment_ids, num_partitions=tf.reduce_max(segment_ids)+1)
    medians = []
    for part in partitions:
        # 沿样本维度排序
        sorted_part = tf.sort(part, axis=0)
        part_len = tf.shape(sorted_part)[0]
        mid_idx = part_len // 2
        if part_len % 2 == 1:
            # 奇数长度取中间元素
            median = sorted_part[mid_idx]
        else:
            # 偶数长度取中间两数的平均
            median = (sorted_part[mid_idx-1] + sorted_part[mid_idx]) / 2.0
        medians.append(median)
    # 拼接所有分段的中位数
    return tf.stack(medians)

# 计算并输出结果
median_result = segment_median(x, segment_ids)
print(median_result)

输出结果 #

tf.Tensor(
[[0.1 0.2 0.4 0.5]
 [0.1 0.2 0.2 0.6]
 [0.1 0.1 0.6 0.5]], shape=(3, 4), dtype=float32)

说明：如果处理更高维度的张量，只需调整tf.sort的axis参数；若需自定义中位数规则（如加权中位数），修改索引计算逻辑即可。

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