TensorFlow自定义Min-Max池化层实现需求与问题求助

以下是完全基于TensorFlow向量化操作的实现，避免了容易出问题的while循环，同时严格满足你提出的按块内最值出现顺序拼接的需求：

实现代码 #

import tensorflow as tf

class MinMaxPoolingLayer(tf.keras.layers.Layer):
    def __init__(self, k, pad_mode='CONSTANT', pad_value=0.0, **kwargs):
        super().__init__(**kwargs)
        self.k = k  # 窗口分块大小
        self.pad_mode = pad_mode  # 序列长度非k整数倍时的补全模式
        self.pad_value = pad_value

    def build(self, input_shape):
        # 自动识别输入是否带通道维度
        self.num_channels = input_shape[-1] if len(input_shape) == 3 else 1
        super().build(input_shape)

    def call(self, inputs):
        # 处理单通道输入，统一为[batch_size, seq_len, channels]格式
        if len(inputs.shape) == 2:
            inputs = tf.expand_dims(inputs, axis=-1)
        
        batch_size, seq_len, channels = inputs.shape
        # 补全序列长度至k的整数倍
        pad_len = (self.k - seq_len % self.k) % self.k
        if pad_len > 0:
            inputs = tf.pad(inputs, [[0,0], [0, pad_len], [0,0]], 
                           mode=self.pad_mode, constant_values=self.pad_value)
            seq_len += pad_len
        
        num_blocks = seq_len // self.k
        # 将输入分块变形：[batch_size, num_blocks, k, channels]
        blocks = tf.reshape(inputs, [batch_size, num_blocks, self.k, channels])
        
        # 计算每个块的最值及对应索引
        min_vals = tf.reduce_min(blocks, axis=2)
        max_vals = tf.reduce_max(blocks, axis=2)
        min_indices = tf.argmin(blocks, axis=2, output_type=tf.int32)
        max_indices = tf.argmax(blocks, axis=2, output_type=tf.int32)
        
        # 判断最值在块中的出现顺序，决定拼接顺序
        min_first = tf.less(min_indices, max_indices)
        # 向量化选择拼接顺序，避免循环
        block_output = tf.where(
            tf.expand_dims(min_first, axis=2),
            tf.concat([tf.expand_dims(min_vals, axis=2), tf.expand_dims(max_vals, axis=2)], axis=2),
            tf.concat([tf.expand_dims(max_vals, axis=2), tf.expand_dims(min_vals, axis=2)], axis=2)
        )
        
        # 整理输出形状，单通道输入自动去掉通道维度
        output = tf.reshape(block_output, [batch_size, num_blocks*2, channels])
        if self.num_channels == 1:
            output = tf.squeeze(output, axis=-1)
        
        return output

    def compute_output_shape(self, input_shape):
        # 提前计算输出形状，兼容Keras模型构建
        if len(input_shape) == 2:
            batch_size, seq_len = input_shape
            pad_len = (self.k - seq_len % self.k) % self.k
            output_seq_len = ((seq_len + pad_len) // self.k) * 2
            return (batch_size, output_seq_len)
        else:
            batch_size, seq_len, channels = input_shape
            pad_len = (self.k - seq_len % self.k) % self.k
            output_seq_len = ((seq_len + pad_len) // self.k) * 2
            return (batch_size, output_seq_len, channels)

测试验证 #

# 用你给出的示例测试
test_input = tf.convert_to_tensor([[1,2,3,6,5,4]], dtype=tf.float32)
layer = MinMaxPoolingLayer(k=3)
output = layer(test_input)
print(output.numpy())
# 输出结果：[[1. 3. 6. 4.]]，完全符合预期

方案优势 #

1. 无循环高效实现：全程使用TensorFlow原生向量化操作，避免了tf.while_loop带来的形状推断问题和性能瓶颈，同时兼容Graph模式和分布式训练。
2. 灵活的输入处理：支持单通道（[batch_size, seq_len]）和多通道（[batch_size, seq_len, channels]）输入，自动适配通道维度。
3. 序列补全支持：当输入序列长度不是k的整数倍时，可通过补全（默认补零）保证分块完整，补全模式和值可自定义。
4. 严格满足需求：通过判断最值在块中的索引位置，确保按原块内出现顺序拼接最值，完全匹配你给出的示例逻辑。

内容的提问来源于stack exchange，提问作者AB Music Box