如何在TensorFlow中基于绝对值大小实现正负值的最大池化？

你需要的是基于绝对值大小选择元素的池化逻辑，而非传统的数值最大池化。核心思路是先找到指定轴上绝对值最大的元素的位置，再从原张量中提取对应元素，具体实现如下：

步骤1：构造输入张量 #

先复现你提供的3D输入张量：

import tensorflow as tf

input_tensor = tf.constant([
    [[2, 100, 1], [7, 4, 1]],
    [[5, 3, 1], [15, 11, 1]],
    [[1000, 323, 1], [-100000, 231, 1]]
], dtype=tf.int32)

步骤2：计算绝对值并获取最大索引 #

通过计算张量的绝对值，找到axis=1维度上绝对值最大的元素的索引：

# 计算每个元素的绝对值
abs_tensor = tf.abs(input_tensor)
# 在axis=1维度上，获取绝对值最大元素的索引
max_abs_indices = tf.argmax(abs_tensor, axis=1, output_type=tf.int32)

步骤3：根据索引提取目标元素 #

调整索引的维度，匹配原张量的结构后，提取对应元素：

# 为索引增加一个维度，满足gather_nd的输入要求
expanded_indices = tf.expand_dims(max_abs_indices, axis=-1)
# 提取原张量中对应索引的元素，batch_dims=1表示保留前1个维度的批量结构
result = tf.gather_nd(input_tensor, expanded_indices, batch_dims=1)

验证结果 #

运行代码后，输出结果与你的预期完全一致：

print(result)

输出：

tf.Tensor(
[[      7     100       1]
 [     15      11       1]
 [-100000      323       1]], shape=(3, 3), dtype=int32)

逻辑说明 #

• 传统tf.reduce_max是比较元素的数值大小，而这里通过先取绝对值，再找最大索引的方式，实现了基于绝对值的最大选择；
• batch_dims=1参数确保批量维度（这里是第0维）不被索引干扰，每个批量样本独立处理；
• 该方法适用于任意维度的张量，只需调整axis参数即可。

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