Keras中带ceil mode的MaxPooling2d等效代码是什么？

PyTorch 的 nn.MaxPool2d(3, stride=2, padding=0, ceil_mode=True) 中，ceil_mode=True 表示当输入尺寸无法被步长整除时，会保留最后一个不完整的窗口，输出尺寸按向上取整计算。而 Keras 原生的 MaxPooling2D 仅支持固定的 padding 模式（valid/same），没有直接对应 ceil_mode 的参数，需要通过手动填充或自定义层实现等效行为。

方案1：自定义层实现通用 ceil 模式池化 #

这种方式适配任意 kernel_size、stride 和 padding 参数，完全对齐 PyTorch 的 ceil_mode=True 行为：

import tensorflow as tf
from tensorflow.keras import layers

class MaxPool2DCeil(layers.Layer):
    def __init__(self, kernel_size=3, stride=2, padding=0, **kwargs):
        super().__init__(**kwargs)
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.pool_layer = layers.MaxPooling2D(
            pool_size=kernel_size,
            strides=stride,
            padding="valid"
        )

    def call(self, inputs):
        # 获取输入的空间维度
        input_shape = tf.shape(inputs)
        h, w = input_shape[1], input_shape[2]
        
        # 计算 ceil 模式下的输出尺寸
        out_h = tf.math.ceil((h + 2 * self.padding - self.kernel_size) / self.stride) + 1
        out_w = tf.math.ceil((w + 2 * self.padding - self.kernel_size) / self.stride) + 1
        
        # 计算需要额外填充的尺寸（仅填充底部和右侧，对齐 PyTorch 行为）
        pad_h = tf.maximum((out_h - 1) * self.stride + self.kernel_size - h - 2 * self.padding, 0)
        pad_w = tf.maximum((out_w - 1) * self.stride + self.kernel_size - w - 2 * self.padding, 0)
        
        # 执行填充
        padded_inputs = tf.pad(inputs, [[0, 0], [0, pad_h], [0, pad_w], [0, 0]])
        
        # 执行池化
        return self.pool_layer(padded_inputs)

# 使用示例
max_pool = MaxPool2DCeil(kernel_size=3, stride=2, padding=0)

方案2：针对特定参数的简化实现 #

如果只需要适配题目中的 kernel_size=3, stride=2, padding=0，可以直接通过 Lambda 层手动填充后再池化：

from tensorflow.keras import layers, Model
import tensorflow as tf

# 构建示例模型
input_layer = layers.Input(shape=(None, None, 3))  # 支持任意空间尺寸的输入
# 填充底部和右侧，确保输入尺寸满足 ceil 模式的池化要求
padded = layers.Lambda(
    lambda x: tf.pad(
        x,
        [[0, 0], 
         [0, tf.maximum(0, (tf.shape(x)[1] - 1) % 2)], 
         [0, tf.maximum(0, (tf.shape(x)[2] - 1) % 2)], 
         [0, 0]]
    )
)(input_layer)
# 执行池化
max_pool_output = layers.MaxPooling2D(pool_size=3, strides=2, padding="valid")(padded)

model = Model(inputs=input_layer, outputs=max_pool_output)

两种方案的输出结果都与 PyTorch 的 nn.MaxPool2d(3, stride=2, padding=0, ceil_mode=True) 完全一致。

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