tf.nn.max_pool_with_argmax 的快速实现与反卷积过程

1. tf.nn.max_pool_with_argmax()：带索引的最大池化，最大池化同时返回索引值

函数原型：
tf.nn.max_pool_with_argmax(
input,
ksize,
strides,
padding,
Targmax=tf.int64,
name=None
)
参数:
input: 一个4-D的Tensor. 它的类型必须为: float32, float64, int32, uint8, int16, int8, int64, bfloat16, uint16, half, uint32, uint64.形状为[batch,height, width, channels].

ksize: 池化核，length >= 4.

strides: 池化的步长length >= 4.

padding: 填充方式，为 "SAME", "VALID".

Targmax: 可选的tf.DType 值类型为: tf.int32, tf.int64. Defaults to tf.int64.

name: 操作名(可选).

返回值:
Tensor对象元组(output, argmax).

output: 一个Tensor.与输入有相同的类型.

argmax: 一个Tensor与Targmax有相同的类型.

2. tf.scatter_nd():根据索引（indices）将给定的张量（updates）散布到新的（初始为零）形状为shape 的张量

函数原型：
tf.scatter_nd(
indices,
updates,
shape,
name=None
)
函数参数：
indices：一个Tensor；必须是以下类型之一：int32，int64；指数张量。
updates：一个Tensor；分散到输出的更新。
shape：一个Tensor；必须与indices具有相同的类型；1-d；得到的张量的形状。
name：操作的名称（可选）。
函数返回值：
此函数将返回一个Tensor，它与updates有相同的类型；根据indices应用的一个新具有给定的形状和更新的张量。

形象化展示

TensorFlow常用函数（四）_蓬莱道人的博客-CSDN博客

tensorflow-1的版本

class MaxUnpooling2D(Layer):
def init(self, size=(2, 2), **kwargs):
super(MaxUnpooling2D, self).init(**kwargs)
self.size = size

def call(self, inputs, output_shape=None):
    updates, mask = inputs[0], inputs[1]
    with K.tf.variable_scope(self.name):
        mask = K.cast(mask, "int32")
        input_shape = K.tf.shape(updates, out_type="int32")
        #  calculation new shape
        if output_shape is None:
            output_shape = (
                input_shape[0],
                input_shape[1] * self.size[0],
                input_shape[2] * self.size[1],
                input_shape[3],
            )
        self.output_shape1 = output_shape

        # calculation indices for batch, height, width and feature maps
        one_like_mask = K.ones_like(mask, dtype="int32")
        batch_shape = K.concatenate([[input_shape[0]], [1], [1], [1]], axis=0)
        batch_range = K.reshape(
            K.tf.range(output_shape[0], dtype="int32"), shape=batch_shape
        )
        b = one_like_mask * batch_range
        y = mask // (output_shape[2] * output_shape[3])
        x = (mask // output_shape[3]) % output_shape[2]
        feature_range = K.tf.range(output_shape[3], dtype="int32")
        f = one_like_mask * feature_range

        # transpose indices & reshape update values to one dimension
        updates_size = K.tf.size(updates)
        indices = K.transpose(K.reshape(K.stack([b, y, x, f]), [4, updates_size]))
        values = K.reshape(updates, [updates_size])
        ret = K.tf.scatter_nd(indices, values, output_shape)
        return ret

def compute_output_shape(self, input_shape):
    mask_shape = input_shape[1]
    return (
        mask_shape[0],
        mask_shape[1] * self.size[0],
        mask_shape[2] * self.size[1],
        mask_shape[3],
    )

tensorflow_2的版本

class MaxUnpooling2D(Layer):
def __init__(self, size=(2, 2), **kwargs):
    super(MaxUnpooling2D, self).__init__(**kwargs)
    self.size = size

def call(self, inputs, output_shape=None):
    updates, mask = inputs[0], inputs[1]
    with tf.compat.v1.variable_scope(self.name):
        mask = K.cast(mask, 'int32')
        input_shape = tf.shape(updates, out_type='int32')
        #print(updates.shape)
        #print(mask.shape)
        if output_shape is None:
            output_shape = (
                input_shape[0],
                input_shape[1] * self.size[0],
                input_shape[2] * self.size[1],
                input_shape[3])

        ret = tf.scatter_nd(K.expand_dims(K.flatten(mask)),
                              K.flatten(updates),
                              [K.prod(output_shape)])

        input_shape = updates.shape
        out_shape = [-1,
                     input_shape[1] * self.size[0],
                     input_shape[2] * self.size[1],
                     input_shape[3]]
    return K.reshape(ret, out_shape)

def get_config(self):
    config = super().get_config().copy()
    config.update({
        'size': self.size
    })
    return config

def compute_output_shape(self, input_shape):
    mask_shape = input_shape[1]
    return (
            mask_shape[0],
            mask_shape[1]*self.size[0],
            mask_shape[2]*self.size[1],
            mask_shape[3]
            )

注意自定的功能函数，需要在get_config中定义变量，以保证网络的回传。

def get_config(self):
    config = super().get_config().copy()
    config.update({
        '属性1': self.属性1，
        '属性2': self.属性2，
    })
    return config

Original: https://blog.csdn.net/chen_gong_ping/article/details/122746884
Author: chen_gong_ping
Title: tf.nn.max_pool_with_argmax 的快速实现与反卷积过程