tensorflow中的自定义中位数池化

import tensorflow as tf

class MedianPooling(tf.keras.layers.Layer):
    def __init__(self, pool_size=(2, 2), strides=None, padding='valid', **kwargs):
        super(MedianPooling, self).__init__(**kwargs)
        self.pool_size = pool_size
        self.strides = strides
        self.padding = padding

    def call(self, inputs):
        return tf.nn.pool(inputs, window_shape=self.pool_size, pooling_type='MEDIAN', strides=self.strides, padding=self.padding)

# 使用自定义中位数池化层
model = tf.keras.models.Sequential()
model.add(MedianPooling(pool_size=(2, 2), strides=(2, 2), padding='valid'))