TensorFlow无循环实现：将张量各点减去对应segment的unsorted_segment_mean均值

其实你离最终结果只差一步——把每个segment的均值根据idx映射回和原始points相同的维度，然后直接做逐元素减法就搞定了，完全不需要循环或者tf.py_func这种非原生操作。

完整代码示例（基于你已有的代码续写） #

import tensorflow as tf

x_index = tf.constant([1, 1, 2, 2])
y_index = tf.constant([1, 1, 3, 4])
points = tf.constant([[0.1, 0.1, 0.1], [0.11, 0.11, 0.11], [0.2, 0.3, 0.1], [0.2, 0.4, 0.1]])
points_x_y_indices = tf.transpose(tf.stack([x_index, y_index]))
uniques, idx = tf.raw_ops.UniqueV2(x=points_x_y_indices, axis=[0], out_idx=tf.dtypes.int32)
n_pillars = int(tf.reduce_max(idx))+1
x_means = tf.math.unsorted_segment_mean(points[:, 0], idx, n_pillars)
y_means = tf.math.unsorted_segment_mean(points[:, 1], idx, n_pillars)
z_means = tf.math.unsorted_segment_mean(points[:, 2], idx, n_pillars)

# 核心步骤：将各维度均值合并，再根据idx映射到每个点的位置
# 1. 把x/y/z均值合并成形状为 [n_pillars, 3] 的张量，每个行对应一个segment的三维均值
pillar_means = tf.stack([x_means, y_means, z_means], axis=1)
# 2. 用tf.gather根据idx取出每个点对应的segment均值，得到和points同形状的 [4, 3] 张量
point_means = tf.gather(pillar_means, idx)
# 3. 逐元素减法完成归一化
normalized_points = points - point_means

# 验证结果（可选）
print(tf.print(normalized_points))

为什么这个方法可行？ #

• tf.gather是TensorFlow的原生向量化操作，它会根据idx中的每个索引值，从pillar_means中取出对应的行。比如你的idx是[0, 0, 1, 2]（前两个点属于同一个segment，后两个各属于不同segment），point_means的每一行就会精准对应points中该行所属segment的均值。
• 这种方式完全贴合TensorFlow的计算图模式，效率远高于循环，也不会有兼容性问题。

另一种等价写法（分开处理各维度） #

如果你不想合并均值张量，也可以分别对每个维度做映射再合并：

x_point_means = tf.gather(x_means, idx)
y_point_means = tf.gather(y_means, idx)
z_point_means = tf.gather(z_means, idx)
point_means = tf.stack([x_point_means, y_point_means, z_point_means], axis=1)
normalized_points = points - point_means

两种写法效果完全一致，选你觉得更易读的就行。

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