wehub-resource-sync
72 lines
3.2 KiB
Python
72 lines
3.2 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name, unused-argument, unused-variable
|
|
"""Gradient of pooling in python"""
|
|
|
|
import numpy as np
|
|
|
|
|
|
def pool_grad_nchw(
|
|
a_np, out_grad_np, pool_size, strides, padding, pool_type, ceil_mode, count_include_pad=True
|
|
):
|
|
"""pool_grad for NCHW layout in python"""
|
|
dtype = a_np.dtype
|
|
n, ic, ih, iw = a_np.shape
|
|
kh, kw = pool_size
|
|
sh, sw = strides
|
|
pt, pl, pb, pr = padding
|
|
|
|
pad_np = np.zeros(shape=(n, ic, ih + pt + pb, iw + pl + pr)).astype(dtype)
|
|
no_zero = (range(n), range(ic), (range(pt, ih + pt)), (range(pl, iw + pl)))
|
|
pad_np[np.ix_(*no_zero)] = a_np
|
|
_, _, oh, ow = out_grad_np.shape
|
|
pool_grad_np = np.zeros(shape=a_np.shape)
|
|
pad_pool_grad_np = np.zeros(shape=pad_np.shape)
|
|
|
|
if pool_type == "avg":
|
|
for i in range(oh):
|
|
for j in range(ow):
|
|
if count_include_pad:
|
|
shape = pad_np[:, :, i * sh : i * sh + kh, j * sw : j * sw + kw].shape
|
|
# this can be different from kh*kw if input size cannot divide stride
|
|
pad_count = shape[2] * shape[3]
|
|
else:
|
|
pad_count = np.sum(
|
|
pad_np[:, :, i * sh : i * sh + kh, j * sw : j * sw + kw] > 0, axis=(2, 3)
|
|
)
|
|
# take the first element, as they are the same across batch and channel
|
|
pad_count = pad_count.ravel()[0]
|
|
pad_pool_grad_np[:, :, i * sh : i * sh + kh, j * sw : j * sw + kw] += out_grad_np[
|
|
:, :, i, j
|
|
].reshape(n, ic, 1, 1) / np.maximum(pad_count, 1)
|
|
elif pool_type == "max":
|
|
for i in range(oh):
|
|
for j in range(ow):
|
|
a_patch = pad_np[:, :, i * sh : i * sh + kh, j * sw : j * sw + kw]
|
|
a_patch = np.reshape(a_patch, (n, ic, -1))
|
|
max_indices = np.argmax(a_patch, axis=2)
|
|
c_idx, n_idx = np.meshgrid(range(ic), range(n), sparse=True)
|
|
h_idx, w_idx = np.unravel_index(max_indices, (kh, kw))
|
|
pad_pool_grad_np[:, :, i * sh : i * sh + kh, j * sw : j * sw + kw][
|
|
n_idx, c_idx, h_idx, w_idx
|
|
] += out_grad_np[n_idx, c_idx, i, j]
|
|
for i in range(pool_grad_np.shape[2]):
|
|
for j in range(pool_grad_np.shape[3]):
|
|
pool_grad_np[:, :, i, j] = pad_pool_grad_np[:, :, i + pt, j + pl]
|
|
|
|
return pool_grad_np
|