convert_gl2tf_maxpool2d.py

import math
import numpy as np
import mxnet as mx
import tensorflow as tf


class GluonModel(mx.gluon.HybridBlock):

    def __init__(self,
                 **kwargs):
        super(GluonModel, self).__init__(**kwargs)

        with self.name_scope():
            self.pool = mx.gluon.nn.MaxPool2D(
                pool_size=2,
                strides=2,
                padding=0)

    def hybrid_forward(self, F, x):
        x = self.pool(x)
        return x


def maxpool2d(x,
              pool_size,
              strides,
              padding=0,
              ceil_mode=False,
              name=None):
    """
    Max pooling operation for two dimensional (spatial) data.

    Parameters:
    ----------
    x : Tensor
        Input tensor.
    pool_size : int or tuple/list of 2 int
        Size of the max pooling windows.
    strides : int or tuple/list of 2 int
        Strides of the pooling.
    padding : int or tuple/list of 2 int, default 0
        Padding value for convolution layer.
    ceil_mode : bool, default False
        When `True`, will use ceil instead of floor to compute the output shape.
    name : str, default 'conv2d'
        Layer name.

    Returns:
    -------
    Tensor
        Resulted tensor.
    """
    if isinstance(padding, int):
        padding = (padding, padding)

    if ceil_mode:
        height = x.shape[2]
        out_height = float(height + 2 * padding[0] - pool_size[0]) / strides[0] + 1.0
        if math.ceil(out_height) > math.floor(out_height):
            padding[0] += 1
        width = x.shape[3]
        out_width = float(width + 2 * padding[1] - pool_size[1]) / strides[1] + 1.0
        if math.ceil(out_width) > math.floor(out_width):
            padding[1] += 1

    if (padding[0] > 0) or (padding[1] > 0):
        x = tf.pad(x, [[0, 0], [0, 0], list(padding), list(padding)], mode="REFLECT")

    x = tf.layers.max_pooling2d(
        inputs=x,
        pool_size=pool_size,
        strides=strides,
        padding='valid',
        data_format='channels_first',
        name=name)

    # if isinstance(pool_size, int):
    #     pool_size = (pool_size, pool_size)
    # if isinstance(strides, int):
    #     strides = (strides, strides)
    # x = tf.nn.max_pool(
    #     value=x,
    #     ksize=(1, 1) + pool_size,
    #     strides=(1, 1) + strides,
    #     padding='VALID',
    #     data_format='NCHW',
    #     name=name)

    return x


def tensorflow_model(x):

    x = maxpool2d(
        x=x,
        pool_size=2,
        strides=2,
        padding=0,
        ceil_mode=False,
        name="pool")
    return x


def main():

    success = True
    for i in range(10):
        x = np.random.randn(10, 10, 224, 224).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)

        gl_x = mx.nd.array(x, ctx)
        gl_y = gl_model(gl_x).asnumpy()

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, 10, 224, 224),
            name='xx')
        tf_model = tensorflow_model(xx)
        with tf.Session() as sess:
            tf_y = sess.run(tf_model, feed_dict={xx: x})
        tf.reset_default_graph()

        dist = np.sum(np.abs(gl_y - tf_y))
        if dist > 1e-5:
            success = False
            print("i={}, dist={}".format(i, dist))
            # print(gl_y)
            # print(tf_y)

    if success:
        print("All ok.")


if __name__ == '__main__':
    main()