Hetu/python/hetu/gpu_ops/MaxPool.py

from __future__ import absolute_import
import numpy as np
from .Node import Op
from .._base import DNNL_LIB
from ..cpu_links import max_pool as cpu_max_pooling
from ..cpu_links import max_pool_gradient as cpu_max_pooling_gradient
from ..gpu_links import CuDNN_max_pooling2d
from ..gpu_links import CuDNN_max_pooling2d_gradient


def np_max_pooling(input, kernel_H, kernel_W, padding=0, stride=1):
    N, C, H, W = input.shape
    assert((H + 2 * padding - kernel_H) % stride == 0)
    assert((W + 2 * padding - kernel_W) % stride == 0)
    pooled_H = (H + 2 * padding - kernel_H) // stride + 1
    pooled_W = (W + 2 * padding - kernel_W) // stride + 1

    pooled_layer = np.zeros(shape=(N, C, pooled_H, pooled_W), dtype=np.float32)
    pooling_size = kernel_H * kernel_W

    for n in range(N):
        for c in range(C):
            for h in range(pooled_H):
                for w in range(pooled_W):
                    hs = h * stride - padding
                    ws = w * stride - padding
                    hend = min(hs + kernel_H, H)
                    wend = min(ws + kernel_W, W)
                    hs = max(hs, 0)
                    ws = max(ws, 0)

                    hargmax = hs
                    wargmax = ws
                    for i in range(hs, hend):
                        for j in range(ws, wend):
                            if input[n][c][i][j] > input[n][c][hargmax][wargmax]:
                                hargmax = i
                                wargmax = j
                    pooled_layer[n][c][h][w] = input[n][c][hargmax][wargmax]

    return pooled_layer


def np_max_pooling_gradient(input, gradient_y, kernel_H, kernel_W, padding=0, stride=1):
    N, C, pooled_H, pooled_W = gradient_y.shape
    H = (pooled_H - 1) * stride + kernel_H - 2 * padding
    W = (pooled_W - 1) * stride + kernel_W - 2 * padding
    gradient_x = np.zeros(shape=(N, C, H, W), dtype=np.float32)
    pooling_size = kernel_H * kernel_W

    for n in range(N):
        for c in range(C):
            for h in range(pooled_H):
                for w in range(pooled_W):
                    hs = h * stride - padding
                    ws = w * stride - padding
                    hend = min(hs + kernel_H, H)
                    wend = min(ws + kernel_W, W)
                    hs = max(hs, 0)
                    ws = max(ws, 0)

                    hargmax = hs
                    wargmax = ws
                    for i in range(hs, hend):
                        for j in range(ws, wend):
                            if input[n][c][i][j] > input[n][c][hargmax][wargmax]:
                                hargmax = i
                                wargmax = j
                    gradient_x[n][c][hargmax][wargmax] += gradient_y[n][c][h][w]

    return gradient_x


class Max_Pool2dOp(Op):
    def __init__(self, node_A, kernel_H, kernel_W, padding, stride, ctx=None):
        super().__init__(Max_Pool2dOp, [node_A], ctx)
        self.padding = padding
        self.stride = stride
        self.kernel_H = kernel_H
        self.kernel_W = kernel_W

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            if DNNL_LIB['DnnlMaxPool']:
                cpu_max_pooling(
                    input_vals[0], self.kernel_H, self.kernel_W, output_val, self.padding, self.stride)
            else:
                output_val[:] = np_max_pooling(input_vals[0].asnumpy(
                ), self.kernel_H, self.kernel_W, self.padding, self.stride)
        else:
            CuDNN_max_pooling2d(
                input_vals[0], self.kernel_H, self.kernel_W, output_val, self.padding, self.stride, stream_handle)

    def gradient(self, output_grad):
        return [max_pool2d_gradient_op(self, output_grad, self.inputs[0], self.kernel_H, self.kernel_W, self.padding, self.stride, ctx=self.raw_ctx)]

    def infer_shape(self, input_shapes):
        """Need to handle input_vals[0].shape != input_vals[1].shape"""
        assert len(input_shapes) == 1
        N, C, H, W = input_shapes[0]
        p_H = (H + 2 * self.padding - self.kernel_H) // self.stride + 1
        p_W = (W + 2 * self.padding - self.kernel_W) // self.stride + 1
        return (N, C, p_H, p_W)


class Max_Pool2d_GradientOp(Op):
    def __init__(self, node_out, node_out_gradient, node_in, kernel_H, kernel_W, padding, stride, ctx=None):
        super().__init__(Max_Pool2d_GradientOp, [
            node_out, node_out_gradient, node_in], ctx)
        self.padding = padding
        self.stride = stride
        self.kernel_H = kernel_H
        self.kernel_W = kernel_W

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            if DNNL_LIB['DnnlMaxPool_Gradient']:
                cpu_max_pooling_gradient(
                    input_vals[2], input_vals[1], self.kernel_H, self.kernel_W, output_val, self.padding, self.stride)
            else:
                output_val[:] = np_max_pooling_gradient(input_vals[2].asnumpy(
                ), input_vals[1].asnumpy(), self.kernel_H, self.kernel_W, self.padding, self.stride)
        else:
            CuDNN_max_pooling2d_gradient(
                input_vals[0], input_vals[1], input_vals[2], self.kernel_H, self.kernel_W, output_val, self.padding, self.stride, stream_handle)

    def gradient(self, output_grad):
        raise NotImplementedError

    def infer_shape(self, input_shapes):
        assert len(input_shapes) == 3
        return input_shapes[2]


def max_pool2d_op(node_A, kernel_H, kernel_W, padding, stride, ctx=None):
    """Make a new instance of Max_Pool2dOp and call the instance.

    Parameters:
    ----
    node_A : Node
        Input Node
    kernel_H : scalar value
        Size of pool(height)
    kernel_W : scalar value
        Size of pool(width)
    padding : scalar value
        Padding edge
    stride : scalar value
        Step Length of the kernel

    Returns:
    ----
    A new Node instance created by Op.

    """
    return Max_Pool2dOp(node_A, kernel_H, kernel_W, padding, stride, ctx=ctx)


def max_pool2d_gradient_op(node_out, node_out_gradient, node_in, kernel_H, kernel_W, padding, stride, ctx=None):
    """Make a new instance of Max_Pool2d_GradientOp and call the instance.

    Parameters:
    ----
    node_out : Node
        Output Node
    node_out_gradient : Node
        Gradient array
    node_in : Node
        Input Node
    kernel_H : scalar value
        Size of pool(height)
    kernel_W : scalar value
        Size of pool(width)
    padding : scalar value
        Padding edge
    stride : scalar value
        Step Length of the kernel

    Returns:
    ----
    A new Node instance created by Op.

    """
    return Max_Pool2d_GradientOp(node_out, node_out_gradient, node_in, kernel_H, kernel_W, padding, stride, ctx=ctx)