wehub-resource-sync
131 lines
4.4 KiB
Python
131 lines
4.4 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
|
|
# ruff: noqa: E741, RUF005
|
|
"""adaptive pool in python"""
|
|
|
|
import numpy as np
|
|
|
|
|
|
def _start_index(index, odim, idim):
|
|
return int(np.floor(index * idim / odim))
|
|
|
|
|
|
def _end_index(index, odim, idim):
|
|
return int(np.ceil((index + 1) * idim / odim))
|
|
|
|
|
|
def _pool1d(in_size, out_size, np_data, np_op):
|
|
out = np.zeros(out_size).astype(np_data.dtype)
|
|
ow = out_size[0]
|
|
for l in range(ow):
|
|
l_start = _start_index(l, ow, in_size[0])
|
|
l_end = _end_index(l, ow, in_size[0])
|
|
l_sl = slice(l_start, l_end)
|
|
out[l] = np_op(np_data[l_sl])
|
|
return out
|
|
|
|
|
|
def _pool2d(in_size, out_size, np_data, np_op):
|
|
out = np.zeros(out_size).astype(np_data.dtype)
|
|
oh, ow = out_size
|
|
for k in range(oh):
|
|
k_start = _start_index(k, oh, in_size[0])
|
|
k_end = _end_index(k, oh, in_size[0])
|
|
k_sl = slice(k_start, k_end)
|
|
for l in range(ow):
|
|
l_start = _start_index(l, ow, in_size[1])
|
|
l_end = _end_index(l, ow, in_size[1])
|
|
l_sl = slice(l_start, l_end)
|
|
out[k, l] = np_op(np_data[k_sl, l_sl])
|
|
return out
|
|
|
|
|
|
def _pool3d(in_size, out_size, np_data, np_op):
|
|
out = np.zeros(out_size).astype(np_data.dtype)
|
|
od, oh, ow = out_size
|
|
for m in range(od):
|
|
m_start = _start_index(m, od, in_size[0])
|
|
m_end = _end_index(m, od, in_size[0])
|
|
m_sl = slice(m_start, m_end)
|
|
for k in range(oh):
|
|
k_start = _start_index(k, oh, in_size[1])
|
|
k_end = _end_index(k, oh, in_size[1])
|
|
k_sl = slice(k_start, k_end)
|
|
for l in range(ow):
|
|
l_start = _start_index(l, ow, in_size[2])
|
|
l_end = _end_index(l, ow, in_size[2])
|
|
l_sl = slice(l_start, l_end)
|
|
out[m, k, l] = np_op(np_data[m_sl, k_sl, l_sl])
|
|
return out
|
|
|
|
|
|
def adaptive_pool_channel_first(np_data, out_size, pool_op, np_op):
|
|
"""The reference function for adaptive pool, channel first layout"""
|
|
ishape = np_data.shape
|
|
n, c = ishape[:2]
|
|
oshape = (n, c) + out_size
|
|
np_out = np.zeros(oshape).astype(np_data.dtype)
|
|
|
|
for i in range(n):
|
|
for j in range(c):
|
|
np_out[i, j] = pool_op(ishape[2:], out_size, np_data[i, j], np_op)
|
|
|
|
return np_out
|
|
|
|
|
|
def adaptive_pool_channel_last(np_data, out_size, pool_op, np_op):
|
|
"""The reference function for adaptive pool, channel last layout"""
|
|
ishape = np_data.shape
|
|
n, c = ishape[0], ishape[-1]
|
|
oshape = (n,) + out_size + (c,)
|
|
np_out = np.zeros(oshape).astype(np_data.dtype)
|
|
|
|
for i in range(n):
|
|
for j in range(c):
|
|
if len(out_size) == 1:
|
|
np_out[i, :, j] = pool_op(ishape[1:-1], out_size, np_data[i, :, j], np_op)
|
|
elif len(out_size) == 2:
|
|
np_out[i, :, :, j] = pool_op(ishape[1:-1], out_size, np_data[i, :, :, j], np_op)
|
|
else:
|
|
np_out[i, :, :, :, j] = pool_op(
|
|
ishape[1:-1], out_size, np_data[i, :, :, :, j], np_op
|
|
)
|
|
|
|
return np_out
|
|
|
|
|
|
def adaptive_pool(np_data, out_size, pool_type, layout):
|
|
"""The reference function for adaptive pool, for 2d and 3d"""
|
|
if isinstance(out_size, int):
|
|
out_size = (out_size,)
|
|
if len(out_size) == 1:
|
|
pool_op = _pool1d
|
|
elif len(out_size) == 2:
|
|
pool_op = _pool2d
|
|
else:
|
|
assert len(out_size) == 3
|
|
pool_op = _pool3d
|
|
|
|
np_op = np.mean if pool_type == "avg" else np.max
|
|
|
|
if layout in ["NCW", "NCHW", "NCDHW"]:
|
|
return adaptive_pool_channel_first(np_data, out_size, pool_op, np_op)
|
|
|
|
assert layout in ["NWC", "NHWC", "NDHWC"]
|
|
return adaptive_pool_channel_last(np_data, out_size, pool_op, np_op)
|