chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,80 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
|
||||
"""TOPI Testing Util functions.
|
||||
|
||||
Used to verify the correctness of operators in TOPI .
|
||||
"""
|
||||
|
||||
from .conv1d_ncw_python import conv1d_ncw_python, group_conv1d_ncw_python
|
||||
from .conv2d_hwcn_python import conv2d_hwcn_python
|
||||
from .conv2d_nchw_python import conv2d_nchw_python
|
||||
from .conv2d_nhwc_python import conv2d_nhwc_python
|
||||
from .conv3d_ncdhw_python import conv3d_ncdhw_python
|
||||
from .conv3d_ndhwc_python import conv3d_ndhwc_python
|
||||
from .conv3d_transpose_ncdhw_python import conv3d_transpose_ncdhw_python
|
||||
from .conv2d_transpose_python import conv2d_transpose_nchw_python, conv2d_transpose_nhwc_python
|
||||
from .conv1d_transpose_ncw_python import (
|
||||
conv1d_transpose_ncw_python,
|
||||
group_conv1d_transpose_ncw_python,
|
||||
)
|
||||
from .correlation_nchw_python import correlation_nchw_python
|
||||
from .deformable_conv2d_python import deformable_conv2d_nchw_python, deformable_conv2d_nhwc_python
|
||||
from .depthwise_conv2d_python import (
|
||||
depthwise_conv2d_python_nchw,
|
||||
depthwise_conv2d_python_nhwc,
|
||||
depthwise_conv2d_python_nchwc,
|
||||
)
|
||||
from .dilate_python import dilate_python
|
||||
from .softmax_python import softmax_python, log_softmax_python
|
||||
from .resize_python import resize1d_python, resize2d_python, resize3d_python
|
||||
from .reorg_python import reorg_python
|
||||
from .roi_align_python import roi_align_nchw_python, roi_align_nhwc_python
|
||||
from .roi_pool_python import roi_pool_nchw_python
|
||||
from .instance_norm_python import instance_norm_python
|
||||
from .layer_norm_python import layer_norm_python
|
||||
from .group_norm_python import group_norm_python
|
||||
from .rms_norm_python import rms_norm_python
|
||||
from .lrn_python import lrn_python
|
||||
from .l2_normalize_python import l2_normalize_python
|
||||
from .gather_python import gather_python
|
||||
from .gather_nd_python import gather_nd_python
|
||||
from .get_valid_counts_python import get_valid_counts_python
|
||||
from .strided_slice_python import strided_slice_python, strided_set_python
|
||||
from .batch_matmul import batch_matmul
|
||||
from .batch_norm import batch_norm
|
||||
from .nms_python import non_max_suppression_python
|
||||
from .slice_axis_python import slice_axis_python
|
||||
from .sequence_mask_python import sequence_mask
|
||||
from .poolnd_python import poolnd_python
|
||||
from .pool_grad_python import pool_grad_nchw
|
||||
from .one_hot import one_hot
|
||||
from .depth_to_space import depth_to_space_python
|
||||
from .space_to_depth import space_to_depth_python
|
||||
from .crop_and_resize_python import crop_and_resize_python
|
||||
from .adaptive_pool_python import adaptive_pool
|
||||
from .grid_sample_python import affine_grid_python, grid_sample_python
|
||||
from .matrix_set_diag import matrix_set_diag
|
||||
from .space_to_batch_nd import space_to_batch_nd_python
|
||||
from .batch_to_space_nd import batch_to_space_nd_python
|
||||
from .nll_loss import nll_loss
|
||||
from .dense import dense
|
||||
from .searchsorted import searchsorted_ref
|
||||
from .conv2d_backcward_weight_python import conv2d_backward_weight_python
|
||||
from .lstm_python import lstm_python
|
||||
from .attention_python import attention_python
|
||||
@@ -0,0 +1,130 @@
|
||||
# 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)
|
||||
@@ -0,0 +1,125 @@
|
||||
# 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.
|
||||
|
||||
"""Attention operator in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .softmax_python import softmax_python
|
||||
|
||||
|
||||
def attention_python(
|
||||
q: np.ndarray,
|
||||
k: np.ndarray,
|
||||
v: np.ndarray,
|
||||
bias: np.ndarray | None,
|
||||
qk_scale: float,
|
||||
causal: str,
|
||||
window_size: int | None = None,
|
||||
layout: str = "BSNH",
|
||||
): # pylint: disable=too-many-arguments, too-many-locals, invalid-name
|
||||
"""Attention operator in python
|
||||
|
||||
Parameters
|
||||
----------
|
||||
q : np.ndarray
|
||||
Query tensor with shape [batch, seq_length, num_heads, head_dim] in the layout specified by
|
||||
`layout`.
|
||||
k : np.ndarray
|
||||
Key tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim] in the layout specified
|
||||
by `layout`.
|
||||
v : np.ndarray
|
||||
Value tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim_v] in the layout
|
||||
specified by `layout`.
|
||||
bias : np.ndarray
|
||||
Bias tensor with shape [batch, num_heads, seq_length, seq_length]
|
||||
qk_scale : float
|
||||
Scale factor for the query-key product.
|
||||
causal : str
|
||||
The type of causal mask to apply. Can be "none", "TopLeft", or "BottomRight".
|
||||
window_size : Optional[int]
|
||||
The window size for the causal mask.
|
||||
layout : str
|
||||
The layout of the input tensors, e.g. "BSNH" or "BNSH".
|
||||
|
||||
Returns
|
||||
-------
|
||||
np.ndarray
|
||||
The output tensor with shape [batch, seq_length, num_heads, head_dim_v] in the layout
|
||||
specified by `layout`.
|
||||
"""
|
||||
assert layout in ["BSNH", "BNSH", "SBNH"]
|
||||
|
||||
dim_b = layout.find("B")
|
||||
dim_s = layout.find("S")
|
||||
dim_n = layout.find("N")
|
||||
dim_h = layout.find("H")
|
||||
|
||||
q = q.transpose(dim_b, dim_n, dim_s, dim_h) # b, n, s, h
|
||||
k = k.transpose(dim_b, dim_n, dim_s, dim_h) # b, n, s_kv, h
|
||||
kt = k.transpose(0, 1, 3, 2) # b, n, h, s_kv
|
||||
v = v.transpose(dim_b, dim_n, dim_s, dim_h)
|
||||
|
||||
num_heads = q.shape[1]
|
||||
num_kv_heads = k.shape[1]
|
||||
s = q.shape[2]
|
||||
s_kv = k.shape[2]
|
||||
|
||||
if num_heads != num_kv_heads:
|
||||
assert num_heads % num_kv_heads == 0
|
||||
factor = num_heads // num_kv_heads
|
||||
kt = np.repeat(kt, factor, axis=1)
|
||||
v = np.repeat(v, factor, axis=1)
|
||||
|
||||
if not qk_scale == "none":
|
||||
score = q @ kt * qk_scale # b, n, s, s_kv
|
||||
else:
|
||||
score = q @ kt / np.sqrt(q.shape[-1]) # b, n, s, s_kv
|
||||
if bias is not None:
|
||||
score = score + bias # b, n, s, s_kv
|
||||
if causal == "none":
|
||||
attn = softmax_python(score, -1)
|
||||
else:
|
||||
if causal == "TopLeft":
|
||||
offset = 0
|
||||
elif causal == "BottomRight":
|
||||
offset = abs(s - s_kv)
|
||||
else:
|
||||
raise ValueError(f"Unsupported causal type: {causal}")
|
||||
score_masked = np.tril(score, k=offset)
|
||||
|
||||
if window_size:
|
||||
score_masked = np.triu(
|
||||
score_masked,
|
||||
-window_size + 1, # pylint: disable=invalid-unary-operand-type
|
||||
)
|
||||
|
||||
score_masked_exp = np.tril(
|
||||
np.exp(score_masked - np.max(score_masked, axis=-1, keepdims=True)), k=offset
|
||||
)
|
||||
|
||||
if window_size:
|
||||
score_masked_exp = np.triu(
|
||||
score_masked_exp,
|
||||
-window_size + 1, # pylint: disable=invalid-unary-operand-type
|
||||
)
|
||||
|
||||
score_masked_sum = np.sum(score_masked_exp, axis=-1, keepdims=True)
|
||||
attn = np.divide(score_masked_exp, score_masked_sum)
|
||||
|
||||
out = attn @ v # b, n, s, h_v
|
||||
return out.transpose(*np.argsort([dim_b, dim_n, dim_s, dim_h]).tolist())
|
||||
@@ -0,0 +1,60 @@
|
||||
# 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
|
||||
"""Batch matmul in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def batch_matmul(x, y, out_dtype=None, trans_x=False, trans_y=True):
|
||||
"""batch_matmul operator implemented in numpy.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : numpy.ndarray
|
||||
3-D with shape [batch, M, K]
|
||||
|
||||
y : numpy.ndarray
|
||||
3-D with shape [batch, N, K]
|
||||
|
||||
out_dtype: string, optional
|
||||
Specify the dtype of output
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : numpy.ndarray
|
||||
3-D with shape [batch, M, N]
|
||||
"""
|
||||
if trans_x:
|
||||
XB, _, M = x.shape
|
||||
else:
|
||||
XB, M, _ = x.shape
|
||||
if trans_y:
|
||||
YB, N, _ = y.shape
|
||||
else:
|
||||
YB, _, N = y.shape
|
||||
batch = max(XB, YB)
|
||||
dtype = x.dtype if out_dtype is None else out_dtype
|
||||
out = np.zeros((batch, M, N)).astype(dtype)
|
||||
for i in range(batch):
|
||||
xx = x[i if XB != 1 else 0].astype(dtype)
|
||||
yy = y[i if YB != 1 else 0].astype(dtype)
|
||||
out[i] = np.dot(
|
||||
xx.T if trans_x else xx,
|
||||
yy.T if trans_y else yy,
|
||||
)
|
||||
return out
|
||||
@@ -0,0 +1,115 @@
|
||||
# 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.
|
||||
"""Batch Normalization implemented in Numpy."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def batch_norm(
|
||||
x: np.ndarray,
|
||||
gamma: np.ndarray,
|
||||
beta: np.ndarray,
|
||||
moving_mean: np.ndarray,
|
||||
moving_var: np.ndarray,
|
||||
axis: int,
|
||||
epsilon: float,
|
||||
center: bool,
|
||||
scale: bool,
|
||||
training: bool,
|
||||
momentum: float,
|
||||
):
|
||||
"""Batch Normalization operator implemented in Numpy.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : np.ndarray
|
||||
Input to be batch-normalized.
|
||||
|
||||
gamma : np.ndarray
|
||||
Scale factor to be applied to the normalized tensor.
|
||||
|
||||
beta : np.ndarray
|
||||
Offset to be applied to the normalized tensor.
|
||||
|
||||
moving_mean : np.ndarray
|
||||
Running mean of input.
|
||||
|
||||
moving_var : np.ndarray
|
||||
Running variance of input.
|
||||
|
||||
axis : int
|
||||
Specify along which shape axis the normalization should occur.
|
||||
|
||||
epsilon : float
|
||||
Small float added to variance to avoid dividing by zero.
|
||||
|
||||
center : bool
|
||||
If True, add offset of beta to normalized tensor, If False,
|
||||
beta is ignored.
|
||||
|
||||
scale : bool
|
||||
If True, scale normalized tensor by gamma. If False, gamma
|
||||
is ignored.
|
||||
|
||||
training : bool
|
||||
Indicating whether it is in training mode. If True, update
|
||||
moving_mean and moving_var.
|
||||
|
||||
momentum : float
|
||||
The value used for the moving_mean and moving_var update
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : np.ndarray
|
||||
Normalized data with same shape as input
|
||||
|
||||
moving_mean : np.ndarray
|
||||
Running mean of input.
|
||||
|
||||
moving_var : np.ndarray
|
||||
Running variance of input.
|
||||
"""
|
||||
shape = [1] * len(x.shape)
|
||||
shape[axis] = x.shape[axis]
|
||||
|
||||
if training:
|
||||
reduce_axes = list(range(len(x.shape)))
|
||||
reduce_axes.remove(axis)
|
||||
reduce_axes = tuple(reduce_axes)
|
||||
data_mean = np.mean(x, axis=reduce_axes)
|
||||
data_var = np.var(x, axis=reduce_axes)
|
||||
data_mean_rs = np.reshape(data_mean, shape)
|
||||
data_var_rs = np.reshape(data_var, shape)
|
||||
out = (x - data_mean_rs) / np.sqrt(data_var_rs + epsilon)
|
||||
else:
|
||||
moving_mean_rs = moving_mean.reshape(shape)
|
||||
moving_var_rs = moving_var.reshape(shape)
|
||||
out = (x - moving_mean_rs) / np.sqrt(moving_var_rs + epsilon)
|
||||
|
||||
if scale:
|
||||
out = out * gamma.reshape(shape)
|
||||
if center:
|
||||
out = out + beta.reshape(shape)
|
||||
|
||||
if training:
|
||||
return [
|
||||
out,
|
||||
(1 - momentum) * moving_mean + momentum * data_mean,
|
||||
(1 - momentum) * moving_var + momentum * data_var,
|
||||
]
|
||||
|
||||
return [out, moving_mean, moving_var]
|
||||
@@ -0,0 +1,99 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Batch to space ND in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from . import strided_slice_python
|
||||
|
||||
|
||||
def batch_to_space_nd_python(data, block_shape, crop_begin_list, crop_end_list):
|
||||
"""Batch to Space operator in python for NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : np.ndarray
|
||||
N-D with shape [batch, spatial_shape, remaining_shapes],
|
||||
where spatial_shape has M dimensions.
|
||||
|
||||
block_shape : list of ints
|
||||
1-D array of size [M] where M is number of spatial dims, specifies block
|
||||
size for each spatial dimension.
|
||||
|
||||
crop_begin_list : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
begin crop size for each spatial dimension.
|
||||
|
||||
crop_end_list : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
end crop size for each spatial dimension.
|
||||
|
||||
Returns
|
||||
-------
|
||||
b2s_out : np.ndarray
|
||||
N-D with shape
|
||||
[batch / prod(block_shape),
|
||||
in_shape[1] * block_shape[0] - crop_begin_list[0] - crop_end_list[0], ...,
|
||||
in_shape[M] * block_shape[M-1] - crop_begin_list[M-1] - crop_end_list[M-1],
|
||||
remaining_shape]
|
||||
"""
|
||||
in_shape = data.shape
|
||||
N = len(in_shape)
|
||||
M = len(block_shape)
|
||||
block_shape_prod = np.prod(block_shape)
|
||||
in_batch = data.shape[0]
|
||||
axis = []
|
||||
r_p_shape = []
|
||||
|
||||
r_shape = [block_shape[i] for i in range(0, M)]
|
||||
axis.append(len(r_shape))
|
||||
r_shape.append(in_batch // block_shape_prod)
|
||||
|
||||
for i in range(1, N):
|
||||
axis.append(len(r_shape))
|
||||
if len(axis) < (M + N):
|
||||
axis.append(len(r_shape) - (M + 1))
|
||||
r_shape.append(in_shape[i])
|
||||
|
||||
r_p_shape.append(int(in_batch / block_shape_prod))
|
||||
for i in range(1, M + 1):
|
||||
r_p_shape.append(in_shape[i] * block_shape[i - 1])
|
||||
for i in range(M + 1, N):
|
||||
r_p_shape.append(in_shape[i])
|
||||
|
||||
b2s_out = np.reshape(data, newshape=r_shape)
|
||||
b2s_out = np.transpose(b2s_out, axes=axis)
|
||||
b2s_out = np.reshape(b2s_out, newshape=r_p_shape)
|
||||
|
||||
# Crop the start and end of dimensions of b2s_out
|
||||
begin_idx = []
|
||||
end_idx = []
|
||||
strides = []
|
||||
|
||||
for i, _ in enumerate(r_p_shape):
|
||||
strides.append(1)
|
||||
if 0 < i <= M:
|
||||
# begin and end index for spatial dimensions
|
||||
begin_idx.append(crop_begin_list[i - 1])
|
||||
end_idx.append(r_p_shape[i] - crop_end_list[i - 1])
|
||||
else:
|
||||
begin_idx.append(0)
|
||||
end_idx.append(r_p_shape[i])
|
||||
|
||||
b2s_out = strided_slice_python(b2s_out, begin_idx, end_idx, strides)
|
||||
return b2s_out
|
||||
@@ -0,0 +1,70 @@
|
||||
# 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
|
||||
"""Common utility for topi test"""
|
||||
|
||||
import numpy as np
|
||||
import scipy.signal
|
||||
|
||||
|
||||
def _convolve2d(data, weights):
|
||||
"""2d convolution operator in HW layout.
|
||||
|
||||
This is intended to be used as a replacement for
|
||||
scipy.signals.convolve2d, with wider support for different dtypes.
|
||||
scipy.signal.convolve2d does not support all TVM-supported
|
||||
dtypes (e.g. float16). Where possible, this function uses
|
||||
scipy.signal.convolve2d to take advantage of compiled scipy
|
||||
routines, falling back to an explicit loop only where needed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
2-D with shape [in_height, in_width]
|
||||
|
||||
weights : numpy.ndarray
|
||||
2-D with shape [filter_height, filter_width].
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
2-D with shape [out_height, out_width]
|
||||
|
||||
Return value and layout conventions are matched to
|
||||
``scipy.signal.convolve2d(data, weights, mode="valid")``
|
||||
"""
|
||||
|
||||
try:
|
||||
return scipy.signal.convolve2d(data, weights, mode="valid")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
weights = np.rot90(weights, k=2)
|
||||
|
||||
assert len(data.shape) == len(weights.shape) == 2
|
||||
|
||||
dtype = data.dtype
|
||||
kernel_h, kernel_w = weights.shape
|
||||
|
||||
output_shape = [a_dim - w_dim + 1 for a_dim, w_dim in zip(data.shape, weights.shape)]
|
||||
output = np.zeros(output_shape, dtype=dtype)
|
||||
|
||||
for y in range(output_shape[0]):
|
||||
for x in range(output_shape[1]):
|
||||
output[y][x] = np.sum(data[y : y + kernel_h, x : x + kernel_w] * weights)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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=unused-variable, invalid-name
|
||||
"""1D convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple1d
|
||||
|
||||
|
||||
def dilate_np(x, dilation):
|
||||
"""1D dilation using numpy
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : numpy.ndarray
|
||||
Array to dilate with shape [batch, in_channel, in_width]
|
||||
|
||||
dilation : int
|
||||
dilation rate of output
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : numpy.ndarray
|
||||
Dilated output with shape [batch, in_channel, (in_width - 1) * dilation + 1]
|
||||
"""
|
||||
irange = range(len(x) - 1)
|
||||
for d in range(dilation - 1):
|
||||
indices = [(d + 1) * (i + 1) for i in irange]
|
||||
x = np.insert(x, indices, 0)
|
||||
return x
|
||||
|
||||
|
||||
def group_conv1d_ncw_python(a_np, w_np, stride, padding, dilation, groups):
|
||||
"Grouped version of `conv1d_ncw_python`, see that for documentation"
|
||||
a_slices = np.array_split(a_np, groups, axis=1)
|
||||
w_slices = np.array_split(w_np, groups, axis=0)
|
||||
b_slices = [
|
||||
conv1d_ncw_python(a_slice, w_slice, stride, padding, dilation)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
return np.concatenate(b_slices, axis=1)
|
||||
|
||||
|
||||
def conv1d_ncw_python(a_np, w_np, stride, padding, dilation):
|
||||
"""1D convolution operator in NCW layout
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
3-D with shape [num_filter, in_channel, filter_width]
|
||||
|
||||
stride : int
|
||||
Stride size
|
||||
|
||||
padding : int, tuple, or str
|
||||
Single int for padding size or tuple of (left, right) padding
|
||||
or a string in ['VALID', 'SAME']
|
||||
|
||||
dilation : int
|
||||
Dilation rate of the kernel
|
||||
|
||||
groups : int
|
||||
Number of groups in the convolution
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : numpy.ndarray
|
||||
3-D with shape [batch, out_channel, out_width]
|
||||
"""
|
||||
batch, in_c, in_w = a_np.shape
|
||||
out_c, _, filter_w = w_np.shape
|
||||
if isinstance(stride, tuple | list):
|
||||
stride = stride[0]
|
||||
if isinstance(dilation, tuple | list):
|
||||
dilation = dilation[0]
|
||||
|
||||
dilated_filter_w = (filter_w - 1) * dilation + 1
|
||||
pad_left, pad_right = get_pad_tuple1d(padding, (dilated_filter_w,))
|
||||
out_w = ((in_w - dilated_filter_w + pad_left + pad_right) // stride) + 1
|
||||
|
||||
padded_a_np = np.zeros((batch, in_c, in_w + pad_left + pad_right))
|
||||
padded_a_np[:, :, pad_left : (in_w + pad_left)] = a_np
|
||||
|
||||
b_np = np.zeros((batch, out_c, out_w))
|
||||
for n in range(batch):
|
||||
for f in range(out_c):
|
||||
for c in range(in_c):
|
||||
out = np.convolve(
|
||||
padded_a_np[n, c], np.flip(dilate_np(w_np[f, c], dilation)), mode="valid"
|
||||
)
|
||||
b_np[n, f] += out[::stride]
|
||||
return b_np
|
||||
@@ -0,0 +1,92 @@
|
||||
# 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=unused-variable
|
||||
"""Transposed 1D convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy
|
||||
|
||||
import tvm.topi.testing
|
||||
from tvm.topi.nn.utils import get_pad_tuple1d
|
||||
|
||||
|
||||
def group_conv1d_transpose_ncw_python(a_np, w_np, stride, padding, output_padding, groups=1):
|
||||
"Grouped version of `conv1d_transpose_ncw_python`, see that for documentation"
|
||||
a_slices = np.array_split(a_np, groups, axis=1)
|
||||
w_slices = np.array_split(w_np, groups, axis=0)
|
||||
b_slices = [
|
||||
conv1d_transpose_ncw_python(a_slice, w_slice, stride, padding, output_padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=1)
|
||||
return b_np
|
||||
|
||||
|
||||
def conv1d_transpose_ncw_python(a_np, w_np, stride, padding, output_padding):
|
||||
"""Transposed 1D convolution operator in NCW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : int or a list/tuple of one int
|
||||
Stride size, or [stride_width]
|
||||
|
||||
padding : int, tuple, or str
|
||||
Single int for padding size, or
|
||||
tuple of 2 ints for left and right padding, or
|
||||
['VALID', 'SAME']
|
||||
|
||||
output_padding : tuple
|
||||
Used to recover the actual output shape in case more than one
|
||||
is possible
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
3-D with shape [batch, out_channel, out_width]
|
||||
|
||||
"""
|
||||
batch, in_c, in_w = a_np.shape
|
||||
_, out_c, filter_w = w_np.shape
|
||||
opad = output_padding[0]
|
||||
if isinstance(stride, int):
|
||||
stride_w = stride
|
||||
else:
|
||||
stride_w = stride[0]
|
||||
assert opad < stride_w
|
||||
fpad_left, fpad_right = get_pad_tuple1d(padding, filter_w)
|
||||
# dilate stage
|
||||
dilated_a_np = tvm.topi.testing.dilate_python(a_np, [1, 1, stride_w])
|
||||
# padding stage
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = filter_w - 1 - fpad_right + opad
|
||||
padded_a_np = np.zeros((batch, in_c, dilated_a_np.shape[2] + bpad_left + bpad_right))
|
||||
padded_a_np[:, :, bpad_left : dilated_a_np.shape[2] + bpad_left] = dilated_a_np
|
||||
# convolution stage
|
||||
out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w + opad
|
||||
b_np = np.zeros((batch, out_c, out_w))
|
||||
for n in range(batch):
|
||||
for f in range(out_c):
|
||||
for c in range(in_c):
|
||||
out = scipy.signal.convolve(padded_a_np[n, c], w_np[c, f], mode="valid")
|
||||
b_np[n, f] += out
|
||||
return b_np
|
||||
@@ -0,0 +1,150 @@
|
||||
# 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, too-many-nested-blocks
|
||||
"""Gradient of conv2d with respect to weight in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
# Reference: cutlass/tools/util/include/cutlass/util/reference/host/convolution.h
|
||||
def conv2d_backward_weight_nchw_python(
|
||||
dy_np, x_np, kernel_size, stride, padding, groups=1, channels=None
|
||||
):
|
||||
"""Gradient of the conv2d op with respect to weight, in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dy_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, out_height, out_width]
|
||||
|
||||
x_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
kernel_size : tuple of two ints
|
||||
Height and width of the weight
|
||||
|
||||
stride : tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : tuple of two ints
|
||||
Spatial padding, or [pad_h, pad_w]
|
||||
|
||||
Returns
|
||||
-------
|
||||
dw_np : np.ndarray
|
||||
4-D with shape [num_filter, in_channel, filter_height, filter_width]
|
||||
|
||||
"""
|
||||
N, C, H, W = x_np.shape
|
||||
_, K, P, Q = dy_np.shape
|
||||
R, S = kernel_size
|
||||
pad_h, pad_w = padding
|
||||
stride_h, stride_w = stride
|
||||
is_depth_wise = C == K and C == groups
|
||||
|
||||
if is_depth_wise:
|
||||
assert channels == groups, "Only channel_mult == 1 supported for now."
|
||||
dw = np.zeros((K, 1, R, S)).astype(dy_np.dtype)
|
||||
else:
|
||||
assert groups == 1, "General grouped conv2d not supported for now."
|
||||
dw = np.zeros((K, C, R, S)).astype(dy_np.dtype)
|
||||
|
||||
for k in range(K):
|
||||
for r in range(R):
|
||||
for s in range(S):
|
||||
for c in range(dw.shape[1]):
|
||||
acc = 0
|
||||
for n in range(N):
|
||||
for p in range(P):
|
||||
for q in range(Q):
|
||||
if not is_depth_wise:
|
||||
in_c = c
|
||||
else:
|
||||
in_c = k
|
||||
|
||||
coord = (
|
||||
n,
|
||||
in_c,
|
||||
p * stride_h - pad_h + r,
|
||||
q * stride_w - pad_w + s,
|
||||
)
|
||||
|
||||
if (
|
||||
coord[2] < H
|
||||
and coord[2] >= 0
|
||||
and coord[3] < W
|
||||
and coord[3] >= 0
|
||||
):
|
||||
acc += dy_np[n, k, p, q] * x_np[coord]
|
||||
|
||||
dw[k, c, r, s] = acc
|
||||
|
||||
return dw
|
||||
|
||||
|
||||
def conv2d_backward_weight_python(
|
||||
dy_np, x_np, kernel_size, stride, padding, layout="NCHW", groups=1, channels=None
|
||||
):
|
||||
"""Gradient of the conv2d op with respect to weight, in NCHW or NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dy_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, out_height, out_width] for NCHW layout
|
||||
|
||||
x_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width] for NCHW layout
|
||||
|
||||
kernel_size : tuple of two ints
|
||||
Height and width of the weight
|
||||
|
||||
stride : tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : tuple of two ints
|
||||
Spatial padding, or [pad_h, pad_w]
|
||||
|
||||
layout: string
|
||||
Layout of dy_np and x_np
|
||||
|
||||
groups: int
|
||||
Number of groups for grouped convolution.
|
||||
|
||||
channels : int
|
||||
Number of output channels of this convolution.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dw_np : np.ndarray
|
||||
Tensor of shape [num_filter, in_channel, filter_height, filter_width] for NCHW layout,
|
||||
[num_filter, filter_height, filter_width, in_channel] for NHWC layout.
|
||||
"""
|
||||
if layout == "NCHW":
|
||||
return conv2d_backward_weight_nchw_python(
|
||||
dy_np, x_np, kernel_size, stride, padding, groups, channels
|
||||
)
|
||||
|
||||
dw_np_oihw = conv2d_backward_weight_nchw_python(
|
||||
np.transpose(dy_np, [0, 3, 1, 2]),
|
||||
np.transpose(x_np, [0, 3, 1, 2]),
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
groups,
|
||||
channels,
|
||||
)
|
||||
return np.transpose(dw_np_oihw, [0, 2, 3, 1])
|
||||
@@ -0,0 +1,79 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy.signal
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple
|
||||
|
||||
|
||||
def conv2d_hwcn_python(a_np, w_np, stride, padding):
|
||||
"""Convolution operator in HWCN layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [in_height, in_width, in_channel, batch]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [filter_height, filter_width, in_channel, num_filter]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of 2 or 4 ints
|
||||
Padding size, or ['VALID', 'SAME'], or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 2 ints
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [out_height, out_width, out_channel, batch]
|
||||
"""
|
||||
in_height, in_width, in_channel, batch = a_np.shape
|
||||
kernel_h, kernel_w, _, num_filter = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
|
||||
pad_h = pad_top + pad_bottom
|
||||
pad_w = pad_left + pad_right
|
||||
# compute the output shape
|
||||
out_channel = num_filter
|
||||
out_height = (in_height - kernel_h + pad_h) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + pad_w) // stride_w + 1
|
||||
# change the layout from HWCN to NCHW
|
||||
at = a_np.transpose((3, 2, 0, 1))
|
||||
wt = w_np.transpose((3, 2, 0, 1))
|
||||
bt = np.zeros((batch, out_channel, out_height, out_width))
|
||||
# computation
|
||||
for n in range(batch):
|
||||
for f in range(out_channel):
|
||||
for c in range(in_channel):
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
apad = np.zeros((in_height + pad_h, in_width + pad_w))
|
||||
apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = at[n, c]
|
||||
else:
|
||||
apad = at[n, c]
|
||||
out = scipy.signal.convolve2d(apad, np.rot90(np.rot90(wt[f, c])), mode="valid")
|
||||
bt[n, f] += out[::stride_h, ::stride_w]
|
||||
return bt.transpose((2, 3, 1, 0))
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals, too-many-branches
|
||||
# ruff: noqa: F841
|
||||
"""Convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple
|
||||
|
||||
|
||||
def _conv2d_nchw_python(a_np, w_np, stride, padding):
|
||||
"""Convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [num_filter, in_channel, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of 2 or 4 ints
|
||||
Padding size, or ['VALID', 'SAME'], or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 2 ints
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
batch, in_channel, in_height, in_width = a_np.shape
|
||||
num_filter, _, kernel_h, kernel_w = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
|
||||
pad_h = pad_top + pad_bottom
|
||||
pad_w = pad_left + pad_right
|
||||
# compute the output shape
|
||||
out_channel = num_filter
|
||||
out_height = (in_height - kernel_h + pad_h) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + pad_w) // stride_w + 1
|
||||
b_np = np.zeros((batch, out_channel, out_height, out_width), dtype=a_np.dtype)
|
||||
# computation
|
||||
for n in range(batch):
|
||||
for f in range(out_channel):
|
||||
for c in range(in_channel):
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
apad = np.zeros((in_height + pad_h, in_width + pad_w), dtype=a_np.dtype)
|
||||
apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = a_np[n, c]
|
||||
else:
|
||||
apad = a_np[n, c]
|
||||
|
||||
out = _conv2d_hw(apad, w_np[f, c])
|
||||
b_np[n, f] += out[::stride_h, ::stride_w]
|
||||
return b_np
|
||||
|
||||
|
||||
def _conv2d_hw(apad, w_np_fc):
|
||||
"""2d convolution operator in HW layout.
|
||||
|
||||
This is intended to be used as a subroutine from
|
||||
_conv2d_nchw_python. Using scipy.signal.convolve2d directly does
|
||||
not work for all dtypes (e.g. float16). Where possible, this
|
||||
function uses scipy.signal.convolve2d to take advantage of
|
||||
compiled scipy routines, falling back to an explicit loop only
|
||||
where needed
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
2-D with shape [in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
2-D with shape [filter_height, filter_width].
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
2-D with shape [out_height, out_width]
|
||||
"""
|
||||
|
||||
try:
|
||||
return scipy.signal.convolve2d(apad, np.rot90(np.rot90(w_np_fc)), mode="valid")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
assert len(apad.shape) == len(w_np_fc.shape) == 2
|
||||
|
||||
dtype = apad.dtype
|
||||
in_height, in_width = apad.shape
|
||||
kernel_h, kernel_w = w_np_fc.shape
|
||||
|
||||
output_shape = [a_dim - w_dim + 1 for a_dim, w_dim in zip(apad.shape, w_np_fc.shape)]
|
||||
output = np.zeros(output_shape, dtype=apad.dtype)
|
||||
|
||||
for y in range(output_shape[0]):
|
||||
for x in range(output_shape[1]):
|
||||
output[y][x] = np.sum(apad[y : y + kernel_h, x : x + kernel_w] * w_np_fc)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def conv2d_nchw_python(a_np, w_np, stride, padding, groups=1):
|
||||
"""Convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [num_filter, in_channel // groups, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of 2 or 4 ints
|
||||
Padding size, or ['VALID', 'SAME'], or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 2 ints
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
a_slices = np.array_split(a_np, groups, axis=1)
|
||||
w_slices = np.array_split(w_np, groups, axis=0)
|
||||
b_slices = [
|
||||
_conv2d_nchw_python(a_slice, w_slice, stride, padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=1)
|
||||
return b_np
|
||||
@@ -0,0 +1,116 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy.signal
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple
|
||||
|
||||
|
||||
def _conv2d_nhwc_python(a_np, w_np, stride, padding):
|
||||
"""Convolution operator in NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [filter_height, filter_width, in_channel, num_filter]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of two ints
|
||||
Padding size, or ['VALID', 'SAME'], or [pad_height, pad_width]
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
"""
|
||||
batch, in_height, in_width, in_channel = a_np.shape
|
||||
kernel_h, kernel_w, _, num_filter = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
|
||||
pad_h = pad_top + pad_bottom
|
||||
pad_w = pad_left + pad_right
|
||||
|
||||
# compute the output shape
|
||||
out_channel = num_filter
|
||||
out_height = (in_height - kernel_h + pad_h) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + pad_w) // stride_w + 1
|
||||
# change the layout from NHWC to NCHW
|
||||
at = a_np.transpose((0, 3, 1, 2))
|
||||
wt = w_np.transpose((3, 2, 0, 1))
|
||||
bt = np.zeros((batch, out_channel, out_height, out_width))
|
||||
# computation
|
||||
for n in range(batch):
|
||||
for f in range(out_channel):
|
||||
for c in range(in_channel):
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
apad = np.zeros((in_height + pad_h, in_width + pad_w))
|
||||
apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = at[n, c]
|
||||
else:
|
||||
apad = at[n, c]
|
||||
out = scipy.signal.convolve2d(apad, np.rot90(np.rot90(wt[f, c])), mode="valid")
|
||||
bt[n, f] += out[::stride_h, ::stride_w]
|
||||
return bt.transpose((0, 2, 3, 1))
|
||||
|
||||
|
||||
def conv2d_nhwc_python(a_np, w_np, stride, padding, groups=1):
|
||||
"""Convolution operator in NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [filter_height, filter_width, in_channel // groups, num_filter]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of 2 or 4 ints
|
||||
Padding size, or ['VALID', 'SAME'], or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 2 ints
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
"""
|
||||
|
||||
a_slices = np.array_split(a_np, groups, axis=3)
|
||||
w_slices = np.array_split(w_np, groups, axis=3)
|
||||
b_slices = [
|
||||
_conv2d_nhwc_python(a_slice, w_slice, stride, padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=3)
|
||||
return b_np
|
||||
@@ -0,0 +1,183 @@
|
||||
# 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=unused-variable
|
||||
"""Transposed convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy
|
||||
|
||||
import tvm.topi.testing
|
||||
from tvm.topi.nn.utils import get_pad_tuple
|
||||
|
||||
|
||||
def _conv2d_transpose_nchw_python(a_np, w_np, stride, padding, output_padding):
|
||||
"""Transposed convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [in_channel, num_filter, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
output_padding : int or a list/tuple of two ints
|
||||
Use to disambiguate the output shape.
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
batch, in_c, in_h, in_w = a_np.shape
|
||||
_, out_c, filter_h, filter_w = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
if isinstance(output_padding, int):
|
||||
opad_h = opad_w = output_padding
|
||||
else:
|
||||
opad_h, opad_w = output_padding
|
||||
assert opad_h < stride_h and opad_w < stride_w
|
||||
# dilate stage
|
||||
dilated_a_np = tvm.topi.testing.dilate_python(a_np, [1, 1, stride_h, stride_w])
|
||||
# padding stage
|
||||
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
||||
bpad_top = filter_h - 1 - fpad_top
|
||||
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = filter_w - 1 - fpad_right + opad_w
|
||||
padded_a_np = np.zeros(
|
||||
(
|
||||
batch,
|
||||
in_c,
|
||||
dilated_a_np.shape[2] + bpad_top + bpad_bottom,
|
||||
dilated_a_np.shape[3] + bpad_left + bpad_right,
|
||||
)
|
||||
).astype(a_np.dtype)
|
||||
padded_a_np[
|
||||
:,
|
||||
:,
|
||||
bpad_top : dilated_a_np.shape[2] + bpad_top,
|
||||
bpad_left : dilated_a_np.shape[3] + bpad_left,
|
||||
] = dilated_a_np
|
||||
# convolution stage
|
||||
out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h + opad_h
|
||||
out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w + opad_w
|
||||
b_np = np.zeros((batch, out_c, out_h, out_w)).astype(a_np.dtype)
|
||||
for n in range(batch):
|
||||
for f in range(out_c):
|
||||
for c in range(in_c):
|
||||
out = scipy.signal.convolve2d(padded_a_np[n, c], w_np[c, f], mode="valid")
|
||||
b_np[n, f] += out
|
||||
return b_np
|
||||
|
||||
|
||||
def conv2d_transpose_nhwc_python(
|
||||
a_nhwc, weight, weight_format, stride, padding, output_padding=(0, 0)
|
||||
):
|
||||
"""Transposed convolution operator in NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_nhwc : numpy.ndarray
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
weight : numpy.ndarray
|
||||
4-D in formats HWIO, HWOI, OIHW or IOHW
|
||||
|
||||
weight_format : str
|
||||
['HWIO', 'HWOI', 'OIHW', 'IOHW']
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
assert a_nhwc.ndim == 4, "a_nhwc number of dimensions should be 4"
|
||||
assert weight.ndim == 4, "weight number of dimensions should be 4"
|
||||
|
||||
a_nchw = np.transpose(a_nhwc, (0, 3, 1, 2))
|
||||
|
||||
# conv2d_transpose_nchw_python needs kernel layout to be IOHW
|
||||
if weight_format == "HWIO":
|
||||
w_iohw = np.transpose(weight, (2, 3, 0, 1))
|
||||
elif weight_format == "HWOI":
|
||||
w_iohw = np.transpose(weight, (3, 2, 0, 1))
|
||||
elif weight_format == "OIHW":
|
||||
w_iohw = np.transpose(weight, (1, 0, 2, 3))
|
||||
elif weight_format == "IOHW":
|
||||
w_iohw = weight
|
||||
else:
|
||||
raise ValueError("Valid weight_formats are HWIO, HWOI, OIHW or IOHW")
|
||||
|
||||
res_nchw = conv2d_transpose_nchw_python(
|
||||
a_nchw, w_iohw, stride, padding, output_padding=output_padding
|
||||
)
|
||||
res_nhwc = np.transpose(res_nchw, (0, 2, 3, 1))
|
||||
return res_nhwc
|
||||
|
||||
|
||||
def conv2d_transpose_nchw_python(a_np, w_np, stride, padding, output_padding, groups=1):
|
||||
"""Convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [in_channel, num_filter // groups, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
output_padding : int or a list/tuple of two ints
|
||||
Use to disambiguate the output shape.
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
a_slices = np.array_split(a_np, groups, axis=1)
|
||||
w_slices = np.array_split(w_np, groups, axis=0)
|
||||
b_slices = [
|
||||
_conv2d_transpose_nchw_python(a_slice, w_slice, stride, padding, output_padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=1)
|
||||
return b_np
|
||||
@@ -0,0 +1,97 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals, too-many-branches
|
||||
"""Convolution 3D in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy.signal
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple3d
|
||||
|
||||
|
||||
def _conv3d_ncdhw_python(a_np, w_np, stride, padding):
|
||||
batch, in_channel, in_depth, in_height, in_width = a_np.shape
|
||||
num_filter, _, kernel_d, kernel_h, kernel_w = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_d = stride_h = stride_w = stride
|
||||
else:
|
||||
stride_d, stride_h, stride_w = stride
|
||||
|
||||
pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right = get_pad_tuple3d(
|
||||
padding, (kernel_d, kernel_h, kernel_w)
|
||||
)
|
||||
pad_d = pad_front + pad_back
|
||||
pad_h = pad_top + pad_bottom
|
||||
pad_w = pad_left + pad_right
|
||||
|
||||
# compute the output shape
|
||||
out_channel = num_filter
|
||||
out_depth = (in_depth - kernel_d + pad_d) // stride_d + 1
|
||||
out_height = (in_height - kernel_h + pad_h) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + pad_w) // stride_w + 1
|
||||
b_np = np.zeros((batch, out_channel, out_depth, out_height, out_width))
|
||||
# computation
|
||||
for n in range(batch):
|
||||
for f in range(out_channel):
|
||||
for c in range(in_channel):
|
||||
if pad_d > 0 or pad_h > 0 or pad_w > 0:
|
||||
apad = np.zeros((in_depth + pad_d, in_height + pad_h, in_width + pad_w))
|
||||
apad[
|
||||
pad_front : pad_front + in_depth,
|
||||
pad_top : pad_top + in_height,
|
||||
pad_left : pad_left + in_width,
|
||||
] = a_np[n, c]
|
||||
else:
|
||||
apad = a_np[n, c]
|
||||
out = scipy.signal.convolve(apad, np.flip(w_np[f, c]), mode="valid")
|
||||
b_np[n, f] += out[::stride_d, ::stride_h, ::stride_w]
|
||||
return b_np
|
||||
|
||||
|
||||
def conv3d_ncdhw_python(a_np, w_np, stride, padding, groups=1):
|
||||
"""Convolution operator in NCDHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of three ints
|
||||
Stride size, or [stride_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of three ints
|
||||
Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
a_slices = np.array_split(a_np, groups, axis=1)
|
||||
w_slices = np.array_split(w_np, groups, axis=0)
|
||||
b_slices = [
|
||||
_conv3d_ncdhw_python(a_slice, w_slice, stride, padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=1)
|
||||
return b_np
|
||||
@@ -0,0 +1,122 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Convolution 3D in python"""
|
||||
|
||||
import numpy as np
|
||||
import scipy.signal
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple3d
|
||||
|
||||
|
||||
def _conv3d_ndhwc_python(a_np, w_np, stride, padding):
|
||||
"""Convolution 3D operator in NDHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of three ints
|
||||
Stride size, or [stride_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of three ints
|
||||
Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
batch, in_depth, in_height, in_width, in_channel = a_np.shape
|
||||
kernel_d, kernel_h, kernel_w, _, num_filter = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_d = stride_h = stride_w = stride
|
||||
else:
|
||||
stride_d, stride_h, stride_w = stride
|
||||
|
||||
pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right = get_pad_tuple3d(
|
||||
padding, (kernel_d, kernel_h, kernel_w)
|
||||
)
|
||||
pad_d = pad_front + pad_back
|
||||
pad_h = pad_top + pad_bottom
|
||||
pad_w = pad_left + pad_right
|
||||
# compute the output shape
|
||||
out_channel = num_filter
|
||||
out_depth = (in_depth - kernel_d + pad_d) // stride_d + 1
|
||||
out_height = (in_height - kernel_h + pad_h) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + pad_w) // stride_w + 1
|
||||
# change the layout from NHWC to NCHW
|
||||
at = a_np.transpose((0, 4, 1, 2, 3))
|
||||
wt = w_np.transpose((4, 3, 0, 1, 2))
|
||||
bt = np.zeros((batch, out_channel, out_depth, out_height, out_width), dtype=a_np.dtype)
|
||||
# computation
|
||||
for n in range(batch):
|
||||
for f in range(out_channel):
|
||||
for c in range(in_channel):
|
||||
if pad_d > 0 or pad_h > 0 or pad_w > 0:
|
||||
apad = np.zeros(
|
||||
(in_depth + pad_d, in_height + pad_h, in_width + pad_w), dtype=a_np.dtype
|
||||
)
|
||||
apad[
|
||||
pad_front : pad_front + in_depth,
|
||||
pad_top : pad_top + in_height,
|
||||
pad_left : pad_left + in_width,
|
||||
] = at[n, c]
|
||||
else:
|
||||
apad = at[n, c]
|
||||
out = scipy.signal.convolve(apad, np.flip(wt[f, c]), mode="valid")
|
||||
bt[n, f] += out[::stride_d, ::stride_h, ::stride_w]
|
||||
return bt.transpose((0, 2, 3, 4, 1))
|
||||
|
||||
|
||||
def conv3d_ndhwc_python(a_np, w_np, stride, padding, groups=1):
|
||||
"""Convolution 3D operator in NDHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of three ints
|
||||
Stride size, or [stride_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of three ints
|
||||
Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
a_slices = np.array_split(a_np, groups, axis=4)
|
||||
w_slices = np.array_split(w_np, groups, axis=4)
|
||||
b_slices = [
|
||||
_conv3d_ndhwc_python(a_slice, w_slice, stride, padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=4)
|
||||
return b_np
|
||||
@@ -0,0 +1,145 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals, too-many-branches
|
||||
# ruff: noqa: F841
|
||||
"""Convolution 3D transpose in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm.topi.testing
|
||||
from tvm.topi.nn.utils import get_pad_tuple3d
|
||||
|
||||
|
||||
def _conv3d_transpose_ncdhw_python(a_np, w_np, stride, padding, output_padding):
|
||||
"""Transposed 3d convolution operator in NCDHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size
|
||||
|
||||
output_padding : int or list/tuple of three ints
|
||||
Used to disambiguate output shape.
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
batch, in_c, in_d, in_h, in_w = a_np.shape
|
||||
_, out_c, filter_d, filter_h, filter_w = w_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_d = stride_h = stride_w = stride
|
||||
else:
|
||||
stride_d, stride_h, stride_w = stride
|
||||
if isinstance(output_padding, int):
|
||||
opad_d = opad_h = opad_w = output_padding
|
||||
else:
|
||||
opad_d, opad_h, opad_w = output_padding
|
||||
assert opad_d < stride_d and opad_h < stride_h and opad_w < stride_w
|
||||
|
||||
# dilate stage
|
||||
dilated_a_np = tvm.topi.testing.dilate_python(a_np, [1, 1, stride_d, stride_h, stride_w])
|
||||
|
||||
# padding stage
|
||||
fpad_front, fpad_top, fpad_left, fpad_back, fpad_bottom, fpad_right = get_pad_tuple3d(
|
||||
padding, (filter_d, filter_h, filter_w)
|
||||
)
|
||||
|
||||
bpad_front = filter_d - 1 - fpad_front
|
||||
bpad_back = filter_d - 1 - fpad_back + opad_d
|
||||
bpad_top = filter_h - 1 - fpad_top
|
||||
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = filter_w - 1 - fpad_right + opad_w
|
||||
|
||||
padded_a_np = np.zeros(
|
||||
(
|
||||
batch,
|
||||
in_c,
|
||||
dilated_a_np.shape[2] + bpad_front + bpad_back,
|
||||
dilated_a_np.shape[3] + bpad_top + bpad_bottom,
|
||||
dilated_a_np.shape[4] + bpad_left + bpad_right,
|
||||
)
|
||||
)
|
||||
|
||||
padded_a_np[
|
||||
:,
|
||||
:,
|
||||
bpad_front : dilated_a_np.shape[2] + bpad_front,
|
||||
bpad_top : dilated_a_np.shape[3] + bpad_top,
|
||||
bpad_left : dilated_a_np.shape[4] + bpad_left,
|
||||
] = dilated_a_np
|
||||
|
||||
# convolution stage
|
||||
out_d = (in_d - 1) * stride_d - bpad_front - bpad_back + filter_d
|
||||
out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h
|
||||
out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w
|
||||
|
||||
w_np = np.flip(w_np, axis=[2, 3, 4]).transpose((1, 0, 2, 3, 4))
|
||||
b_np = tvm.topi.testing.conv3d_ncdhw_python(
|
||||
padded_a_np, w_np, stride=(1, 1, 1), padding=(0, 0, 0)
|
||||
)
|
||||
|
||||
return b_np
|
||||
|
||||
|
||||
def conv3d_transpose_ncdhw_python(a_np, w_np, stride, padding, output_padding, groups=1):
|
||||
"""Transposed 3d convolution operator in NCDHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size
|
||||
|
||||
output_padding : int or list/tuple of three ints
|
||||
Used to disambiguate output shape.
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
a_slices = np.array_split(a_np, groups, axis=1)
|
||||
w_slices = np.array_split(w_np, groups, axis=0)
|
||||
b_slices = [
|
||||
_conv3d_transpose_ncdhw_python(a_slice, w_slice, stride, padding, output_padding)
|
||||
for a_slice, w_slice in zip(a_slices, w_slices)
|
||||
]
|
||||
b_np = np.concatenate(b_slices, axis=1)
|
||||
return b_np
|
||||
@@ -0,0 +1,109 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
# ruff: noqa: E731
|
||||
"""Convolution 3D in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def correlation_nchw_python(
|
||||
data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply
|
||||
):
|
||||
"""Correlationn operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data1_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
data2_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
kernel_size: int
|
||||
Kernel size for correlation, must be an odd number
|
||||
|
||||
max_displacement: int
|
||||
Max displacement of Correlation
|
||||
|
||||
stride1: int
|
||||
Stride for data1
|
||||
|
||||
stride2: int
|
||||
Stride for data2 within the neightborhood centered around data1
|
||||
|
||||
padding: int
|
||||
Padding for correlation
|
||||
|
||||
is_multiply: bool
|
||||
operation type is either multiplication or substraction
|
||||
|
||||
Returns
|
||||
-------
|
||||
c_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
# compute output's dimension
|
||||
pad_data_height = data1.shape[2] + 2 * padding
|
||||
pad_data_width = data1.shape[3] + 2 * padding
|
||||
kernel_radius = (kernel_size - 1) // 2
|
||||
border_size = max_displacement + kernel_radius
|
||||
out_width = (pad_data_width - border_size * 2) // stride1
|
||||
out_height = (pad_data_height - border_size * 2) // stride1
|
||||
neighborhood_grid_radius = max_displacement // stride2
|
||||
neighborhood_grid_width = neighborhood_grid_radius * 2 + 1
|
||||
out_channel = neighborhood_grid_width * neighborhood_grid_width
|
||||
|
||||
out = np.zeros((data1.shape[0], out_channel, out_height, out_width))
|
||||
pad_data1 = np.zeros((data1.shape[0], data1.shape[1], pad_data_height, pad_data_width))
|
||||
pad_data2 = np.zeros((data1.shape[0], data1.shape[1], pad_data_height, pad_data_width))
|
||||
|
||||
pad_data1[:, :, padding : padding + data1.shape[2], padding : padding + data1.shape[3]] = data1[
|
||||
:, :, :, :
|
||||
]
|
||||
pad_data2[:, :, padding : padding + data2.shape[2], padding : padding + data2.shape[3]] = data2[
|
||||
:, :, :, :
|
||||
]
|
||||
|
||||
if is_multiply:
|
||||
corr_func = lambda x, y: x * y
|
||||
else:
|
||||
corr_func = lambda x, y: abs(x - y)
|
||||
|
||||
# pylint: disable=too-many-nested-blocks
|
||||
for i in range(out_height):
|
||||
for j in range(out_width):
|
||||
for nbatch in range(data1.shape[0]):
|
||||
# x1,y1 is the location in data1 , i,j is the location in output
|
||||
x1 = j * stride1 + max_displacement
|
||||
y1 = i * stride1 + max_displacement
|
||||
|
||||
for q in range(out_channel):
|
||||
# location in data2
|
||||
x2 = x1 + (q % neighborhood_grid_width - neighborhood_grid_radius) * stride2
|
||||
y2 = y1 + (q // neighborhood_grid_width - neighborhood_grid_radius) * stride2
|
||||
|
||||
for h in range(kernel_size):
|
||||
for w in range(kernel_size):
|
||||
for channel in range(data1.shape[1]):
|
||||
out[nbatch, q, i, j] += corr_func(
|
||||
pad_data1[nbatch, channel, y1 + h, x1 + w],
|
||||
pad_data2[nbatch, channel, y2 + h, x2 + w],
|
||||
)
|
||||
|
||||
out /= float(kernel_size**2 * data1.shape[1])
|
||||
return out
|
||||
@@ -0,0 +1,121 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals, too-many-nested-blocks
|
||||
"""crop and resize in python"""
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def crop_and_resize_python(
|
||||
image, boxes, box_indices, crop_size, layout, method="bilinear", extrapolation_value=0
|
||||
):
|
||||
"""Crop and resize using python"""
|
||||
(target_h, target_w) = crop_size
|
||||
|
||||
if layout == "NHWC":
|
||||
batch = boxes.shape[0]
|
||||
image_height, image_width, channel = image.shape[1], image.shape[2], image.shape[3]
|
||||
scaled_image = np.ones((batch, target_h, target_w, channel))
|
||||
else:
|
||||
batch = boxes.shape[0]
|
||||
channel, image_height, image_width = image.shape[1], image.shape[2], image.shape[3]
|
||||
scaled_image = np.ones((batch, channel, target_h, target_w))
|
||||
|
||||
for n, box in enumerate(boxes):
|
||||
b_in = box_indices[n]
|
||||
y1, x1 = boxes[n][0], boxes[n][1]
|
||||
y2, x2 = boxes[n][2], boxes[n][3]
|
||||
|
||||
in_h = (image_height - 1) * (y2 - y1)
|
||||
in_w = (image_width - 1) * (x2 - x1)
|
||||
h_scale = np.float32(in_h) / np.float32(target_h - 1)
|
||||
w_scale = np.float32(in_w) / np.float32(target_w - 1)
|
||||
|
||||
for y in range(target_h):
|
||||
in_y = y1 * (image_height - 1) + h_scale * y
|
||||
|
||||
if in_y < 0 or in_y > image_height - 1:
|
||||
for x in range(target_w):
|
||||
for d in range(channel):
|
||||
if layout == "NHWC":
|
||||
scaled_image[n][y][x][d] = extrapolation_value
|
||||
else:
|
||||
scaled_image[n][d][y][x] = extrapolation_value
|
||||
continue
|
||||
|
||||
if method == "bilinear":
|
||||
top_y_index = math.floor(in_y)
|
||||
bottom_y_index = math.ceil(in_y)
|
||||
y_lerp = in_y - top_y_index
|
||||
|
||||
for x in range(target_w):
|
||||
in_x = x1 * (image_width - 1) + x * w_scale
|
||||
if in_x < 0 or in_x > image_width - 1:
|
||||
for d in range(channel):
|
||||
if layout == "NHWC":
|
||||
scaled_image[n][y][x][d] = extrapolation_value
|
||||
else:
|
||||
scaled_image[n][d][y][x] = extrapolation_value
|
||||
continue
|
||||
|
||||
left_x_index = math.floor(in_x)
|
||||
right_x_index = math.ceil(in_x)
|
||||
x_lerp = in_x - left_x_index
|
||||
|
||||
for d in range(channel):
|
||||
if layout == "NHWC":
|
||||
top_left = image[b_in][top_y_index][left_x_index][d]
|
||||
top_right = image[b_in][top_y_index][right_x_index][d]
|
||||
bottom_left = image[b_in][bottom_y_index][left_x_index][d]
|
||||
bottom_right = image[b_in][bottom_y_index][right_x_index][d]
|
||||
top = top_left + (top_right - top_left) * x_lerp
|
||||
bottom = bottom_left + (bottom_right - bottom_left) * x_lerp
|
||||
scaled_image[n][y][x][d] = top + (bottom - top) * y_lerp
|
||||
else:
|
||||
top_left = image[b_in][d][top_y_index][left_x_index]
|
||||
top_right = image[b_in][d][top_y_index][right_x_index]
|
||||
bottom_left = image[b_in][d][bottom_y_index][left_x_index]
|
||||
bottom_right = image[b_in][d][bottom_y_index][right_x_index]
|
||||
top = top_left + (top_right - top_left) * x_lerp
|
||||
bottom = bottom_left + (bottom_right - bottom_left) * x_lerp
|
||||
scaled_image[n][d][y][x] = top + (bottom - top) * y_lerp
|
||||
|
||||
elif method == "nearest_neighbor":
|
||||
for x in range(target_w):
|
||||
in_x = x1 * (image_width - 1) + x * w_scale
|
||||
if in_x < 0 or in_x > image_width - 1:
|
||||
for d in range(channel):
|
||||
if layout == "NHWC":
|
||||
scaled_image[n][y][x][d] = extrapolation_value
|
||||
else:
|
||||
scaled_image[n][d][y][x] = extrapolation_value
|
||||
continue
|
||||
closest_x_index = np.round(in_x).astype("int32")
|
||||
closest_y_index = np.round(in_y).astype("int32")
|
||||
for d in range(channel):
|
||||
if layout == "NHWC":
|
||||
scaled_image[n][y][x][d] = image[b_in][closest_y_index][
|
||||
closest_x_index
|
||||
][d]
|
||||
else:
|
||||
scaled_image[n][d][y][x] = image[b_in][d][closest_y_index][
|
||||
closest_x_index
|
||||
]
|
||||
|
||||
return scaled_image
|
||||
@@ -0,0 +1,181 @@
|
||||
# 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, too-many-locals, too-many-arguments
|
||||
"""Deformable convolution in python"""
|
||||
|
||||
import itertools
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple
|
||||
|
||||
|
||||
def deformable_conv2d_nchw_python(
|
||||
a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups
|
||||
):
|
||||
"""Deformable convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
offset_np : numpy.ndarray
|
||||
4-D with shape [batch, deformable_groups * filter_height * filter_width * 2,
|
||||
out_height, out_width]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [num_filter, in_channel, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of 2 or 4 ints
|
||||
Padding size, or ['VALID', 'SAME'], or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 2 ints
|
||||
|
||||
dilation : int or a list/tuple of two ints
|
||||
Dilation size, or [dilate_height, dilate_width]
|
||||
|
||||
deformable_groups : int
|
||||
Number of deformable groups
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
batch, in_channel, in_height, in_width = a_np.shape
|
||||
out_channel, _, kernel_h, kernel_w = w_np.shape
|
||||
out_height, out_width = offset_np.shape[-2:]
|
||||
dtype = a_np.dtype
|
||||
ic_per_dgroup = in_channel // deformable_groups
|
||||
assert groups == 1, "deformable_conv2d_nchw_python does not support groups > 1"
|
||||
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
pad_top, pad_left, _, _ = get_pad_tuple(padding, (kernel_h, kernel_w))
|
||||
|
||||
if isinstance(dilation, int):
|
||||
dilation_h = dilation_w = dilation
|
||||
else:
|
||||
dilation_h, dilation_w = dilation
|
||||
|
||||
def _bilinear(n, c, h, w):
|
||||
y_low = math.floor(h)
|
||||
x_low = math.floor(w)
|
||||
y_high = y_low + 1
|
||||
x_high = x_low + 1
|
||||
|
||||
wy_h = h - y_low
|
||||
wx_h = w - x_low
|
||||
wy_l = 1 - wy_h
|
||||
wx_l = 1 - wx_h
|
||||
|
||||
val = 0
|
||||
for wx, xp in zip((wx_l, wx_h), (x_low, x_high)):
|
||||
for wy, yp in zip((wy_l, wy_h), (y_low, y_high)):
|
||||
if 0 <= yp < in_height and 0 <= xp < in_width:
|
||||
val += wx * wy * a_np[n, c, yp, xp]
|
||||
return val
|
||||
|
||||
a_deform = np.zeros((batch, in_channel, out_height, out_width, kernel_h, kernel_w), dtype=dtype)
|
||||
for n, h, w in itertools.product(range(batch), range(out_height), range(out_width)):
|
||||
offset = offset_np[n, :, h, w].reshape(deformable_groups, kernel_h, kernel_w, 2)
|
||||
in_h = h * stride_h - pad_top
|
||||
in_w = w * stride_w - pad_left
|
||||
|
||||
index_h_base, index_w_base = np.meshgrid(
|
||||
np.arange(in_h, in_h + kernel_h * dilation_h, dilation_h, dtype=offset_np.dtype),
|
||||
np.arange(in_w, in_w + kernel_w * dilation_w, dilation_w, dtype=offset_np.dtype),
|
||||
indexing="ij",
|
||||
)
|
||||
|
||||
for c, kh, kw in itertools.product(range(in_channel), range(kernel_h), range(kernel_w)):
|
||||
dg = c // ic_per_dgroup
|
||||
index_h = index_h_base + offset[dg, ..., 0]
|
||||
index_w = index_w_base + offset[dg, ..., 1]
|
||||
|
||||
y, x = index_h[kh, kw], index_w[kh, kw]
|
||||
if y < 0 or y >= in_height or x < 0 or x >= in_width:
|
||||
continue
|
||||
a_deform[n, c, h, w, kh, kw] = _bilinear(n, c, y, x)
|
||||
|
||||
b_np = np.zeros((batch, out_channel, out_height, out_width), dtype=dtype)
|
||||
for n, c, f, h, w in itertools.product(
|
||||
range(batch), range(in_channel), range(out_channel), range(out_height), range(out_width)
|
||||
):
|
||||
b_np[n, f, h, w] += np.tensordot(a_deform[n, c, h, w], w_np[f, c])
|
||||
|
||||
return b_np
|
||||
|
||||
|
||||
def deformable_conv2d_nhwc_python(
|
||||
a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups
|
||||
):
|
||||
"""Deformable convolution operator in NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
offset_np : numpy.ndarray
|
||||
4-D with shape [batch, out_height, out_width,
|
||||
deformable_groups * filter_height * filter_width * 2]
|
||||
|
||||
w_np : numpy.ndarray
|
||||
4-D with shape [filter_height, filter_width, in_channel, num_filter]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or str or a list/tuple of 2 or 4 ints
|
||||
Padding size, or ['VALID', 'SAME'], or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 2 ints
|
||||
|
||||
dilation : int or a list/tuple of two ints
|
||||
Dilation size, or [dilate_height, dilate_width]
|
||||
|
||||
deformable_groups : int
|
||||
Number of deformable groups
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
a_np = np.transpose(a_np, [0, 3, 1, 2]) # NHWC -> NCHW
|
||||
offset_np = np.transpose(offset_np, [0, 3, 1, 2]) # NHWC -> NCHW
|
||||
w_np = np.transpose(w_np, [3, 2, 0, 1]) # HWIO -> OIHW
|
||||
b_np = deformable_conv2d_nchw_python(
|
||||
a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups
|
||||
)
|
||||
b_np = np.transpose(b_np, [0, 2, 3, 1]) # NCHW -> NHWC
|
||||
return b_np
|
||||
@@ -0,0 +1,54 @@
|
||||
# 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
|
||||
"""Dense in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def dense(x, y, bias, use_bias=False, use_relu=False, out_dtype=None):
|
||||
"""dense operator implemented in numpy.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : numpy.ndarray
|
||||
2-D with shape [M, K]
|
||||
|
||||
y : numpy.ndarray
|
||||
2-D with shape [N, K]
|
||||
|
||||
bias: numpy.ndarray
|
||||
1-D with shape [M,]
|
||||
|
||||
out_dtype: string, optional
|
||||
Specify the dtype of output
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : numpy.ndarray
|
||||
2-D with shape [M, N]
|
||||
"""
|
||||
dtype = x.dtype if out_dtype is None else out_dtype
|
||||
if use_bias:
|
||||
out = np.dot(x.astype(dtype), y.T.astype(dtype)) + bias
|
||||
else:
|
||||
out = np.dot(x.astype(dtype), y.T.astype(dtype))
|
||||
|
||||
if use_relu:
|
||||
out = np.maximum(out, 0)
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,53 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Depth to space in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def depth_to_space_python(data, block_size, mode="DCR"):
|
||||
"""Depth to Space operator in python for NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : np.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
block_size : int
|
||||
Size of blocks to convert channel pixels into.
|
||||
|
||||
Returns
|
||||
-------
|
||||
d2s_out : np.ndarray
|
||||
4-D with shape [batch, in_channel / (block_size * block_size),
|
||||
out_height * block_size, out_width * block_size]
|
||||
"""
|
||||
in_n, in_c, in_h, in_w = data.shape
|
||||
new_h = int(in_h * block_size)
|
||||
new_w = int(in_h * block_size)
|
||||
new_c = int(in_c / (block_size * block_size))
|
||||
|
||||
if mode == "DCR":
|
||||
expanded = np.reshape(data, newshape=[in_n, block_size, block_size, new_c, in_h, in_w])
|
||||
transposed = np.transpose(expanded, axes=[0, 3, 4, 1, 5, 2])
|
||||
else:
|
||||
expanded = np.reshape(data, newshape=(in_n, new_c, block_size, block_size, in_h, in_w))
|
||||
transposed = np.transpose(expanded, axes=(0, 1, 4, 2, 5, 3))
|
||||
newshape = [in_n, new_c, new_h, new_w]
|
||||
d2s_out = np.reshape(transposed, newshape=newshape)
|
||||
return d2s_out
|
||||
@@ -0,0 +1,167 @@
|
||||
# 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-variable, line-too-long
|
||||
"""Depthwise convolution in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.topi.nn.utils import get_pad_tuple
|
||||
|
||||
from .common import _convolve2d
|
||||
|
||||
|
||||
def depthwise_conv2d_python_nchw(input_np, filter_np, stride, padding):
|
||||
"""Depthwise convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
filter_np : numpy.ndarray
|
||||
4-D with shape [in_channel, channel_multiplier, filter_height, filter_width]
|
||||
|
||||
stride : list / tuple of 2 ints
|
||||
[stride_height, stride_width]
|
||||
|
||||
padding : str
|
||||
'VALID' or 'SAME'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
batch, in_channel, in_height, in_width = input_np.shape
|
||||
_, channel_multiplier, filter_height, filter_width = filter_np.shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (filter_height, filter_width))
|
||||
pad_h = pad_top + pad_bottom
|
||||
pad_w = pad_left + pad_right
|
||||
|
||||
out_channel = in_channel * channel_multiplier
|
||||
out_height = (in_height - filter_height + pad_h) // stride_h + 1
|
||||
out_width = (in_width - filter_width + pad_w) // stride_w + 1
|
||||
output_np = np.zeros((batch, out_channel, out_height, out_width))
|
||||
|
||||
for i in range(batch):
|
||||
for j in range(out_channel):
|
||||
apad = input_np[i, j // channel_multiplier, :, :]
|
||||
if pad_h or pad_w:
|
||||
apad = np.pad(apad, [(pad_top, pad_bottom), (pad_left, pad_right)], "constant")
|
||||
|
||||
conv = _convolve2d(
|
||||
apad,
|
||||
np.rot90(filter_np[j // channel_multiplier, j % channel_multiplier, :, :], k=2),
|
||||
)
|
||||
output_np[i, j, :, :] = conv[
|
||||
::stride_h,
|
||||
::stride_w,
|
||||
]
|
||||
|
||||
return output_np
|
||||
|
||||
|
||||
def depthwise_conv2d_python_nchwc(input_np, filter_np, stride, padding):
|
||||
"""Depthwise convolution operator in NCHWc layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_np : numpy.ndarray
|
||||
5-D with shape [batch, in_channel_chunk, in_height, in_width, in_channel_block]
|
||||
|
||||
filter_np : numpy.ndarray
|
||||
6-D with shape [out_channel_chunk, channel_multiplier_chunk,
|
||||
filter_height, filter_width,
|
||||
channel_multiplier_block, out_channel_block]
|
||||
|
||||
stride : list / tuple of 2 ints
|
||||
[stride_height, stride_width]
|
||||
|
||||
padding : str
|
||||
'VALID' or 'SAME'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output_np : np.ndarray
|
||||
5-D with shape [batch, out_channel_chunk, out_height, out_width, out_channel_block]
|
||||
"""
|
||||
# Transform to NCHW
|
||||
batch_size, in_channel_chunk, in_height, in_width, in_channel_block = input_np.shape
|
||||
input_nchw = input_np.transpose(0, 1, 4, 2, 3).reshape(
|
||||
(batch_size, in_channel_chunk * in_channel_block, in_height, in_width)
|
||||
)
|
||||
|
||||
(
|
||||
out_channel_chunk,
|
||||
channel_multiplier_chunk,
|
||||
filter_height,
|
||||
filter_width,
|
||||
channel_multiplier_block,
|
||||
out_channel_block,
|
||||
) = filter_np.shape
|
||||
filter_nchw = filter_np.transpose(0, 5, 1, 4, 2, 3).reshape(
|
||||
(
|
||||
out_channel_chunk * out_channel_block,
|
||||
channel_multiplier_chunk * channel_multiplier_block,
|
||||
filter_height,
|
||||
filter_width,
|
||||
)
|
||||
)
|
||||
|
||||
# Perform conv2d
|
||||
output_np = depthwise_conv2d_python_nchw(input_nchw, filter_nchw, stride, padding)
|
||||
|
||||
# Transform back to NCHWc
|
||||
|
||||
# pylint: disable=unpacking-non-sequence
|
||||
batch_size, out_channel, out_height, out_width = output_np.shape
|
||||
return output_np.reshape(
|
||||
(batch_size, out_channel_chunk, out_channel_block, out_height, out_width)
|
||||
).transpose(0, 1, 3, 4, 2)
|
||||
|
||||
|
||||
def depthwise_conv2d_python_nhwc(input_np, filter_np, stride, padding):
|
||||
"""Depthwise convolution operator in nhwc layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_np : numpy.ndarray
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
filter_np : numpy.ndarray
|
||||
4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
|
||||
|
||||
stride : list / tuple of 2 ints
|
||||
[stride_height, stride_width]
|
||||
|
||||
padding : str
|
||||
'VALID' or 'SAME'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output_np : np.ndarray
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
"""
|
||||
input_nchw = input_np.transpose(0, 3, 1, 2)
|
||||
filter_nchw = filter_np.transpose(2, 3, 0, 1)
|
||||
output_nchw = depthwise_conv2d_python_nchw(input_nchw, filter_nchw, stride, padding)
|
||||
return output_nchw.transpose(0, 2, 3, 1)
|
||||
@@ -0,0 +1,64 @@
|
||||
# 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
|
||||
"""Dilate operation in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def dilate_python(input_np, strides, dilation_value=0.0, out_dtype=None):
|
||||
"""Dilate operation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_np : numpy.ndarray
|
||||
n-D, can be any layout.
|
||||
|
||||
strides : list / tuple of n ints
|
||||
Dilation stride on each dimension, 1 means no dilation.
|
||||
|
||||
dilation_value : int/float, optional
|
||||
Value used to dilate the input.
|
||||
|
||||
out_dtype : Option[str]
|
||||
The datatype of the dilated array. If unspecified, will use
|
||||
the same dtype as the input array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output_np : numpy.ndarray
|
||||
n-D, the same layout as Input.
|
||||
|
||||
"""
|
||||
assert len(input_np.shape) == len(strides), (
|
||||
f"Input dimension and strides size dismatch : {len(input_np.shape)} vs {len(strides)}"
|
||||
)
|
||||
|
||||
if out_dtype is None:
|
||||
out_dtype = input_np.dtype
|
||||
|
||||
output_size = [
|
||||
(input_dim - 1) * stride + 1 for input_dim, stride in zip(input_np.shape, strides)
|
||||
]
|
||||
non_zero_elements = np.ix_(
|
||||
*[range(0, output_dim, stride) for output_dim, stride in zip(output_size, strides)]
|
||||
)
|
||||
|
||||
output_np = np.full(shape=output_size, fill_value=dilation_value, dtype=out_dtype)
|
||||
output_np[non_zero_elements] = input_np
|
||||
|
||||
return output_np
|
||||
@@ -0,0 +1,55 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
# ruff: noqa: RUF005
|
||||
"""gather_nd in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def gather_nd_python(a_np, indices_np):
|
||||
"""Python version of GatherND operator
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
Numpy array
|
||||
|
||||
indices_np : numpy.ndarray
|
||||
Numpy array
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : numpy.ndarray
|
||||
Numpy array
|
||||
"""
|
||||
a_shape = a_np.shape
|
||||
indices_np = indices_np.astype("int32")
|
||||
indices_shape = indices_np.shape
|
||||
assert len(indices_shape) > 1
|
||||
assert indices_shape[0] <= len(a_shape)
|
||||
b_shape = list(indices_shape[1:])
|
||||
for i in range(indices_shape[0], len(a_shape)):
|
||||
b_shape.append(a_shape[i])
|
||||
b_np = np.zeros(b_shape)
|
||||
for idx in np.ndindex(*indices_shape[1:]):
|
||||
a_idx = []
|
||||
for i in range(indices_shape[0]):
|
||||
indices_pos = tuple([i] + list(idx))
|
||||
a_idx.append(indices_np[indices_pos])
|
||||
b_np[idx] = a_np[tuple(a_idx)]
|
||||
return b_np
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""gather in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def gather_python(data, axis, indices):
|
||||
"""Python version of Gather operator
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
Numpy array
|
||||
|
||||
axis: int
|
||||
integer
|
||||
|
||||
indices : numpy.ndarray
|
||||
Numpy array
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : numpy.ndarray
|
||||
Numpy array
|
||||
"""
|
||||
shape_indices = indices.shape
|
||||
out = np.zeros(shape_indices, dtype=data.dtype)
|
||||
for index in np.ndindex(*shape_indices):
|
||||
new_index = list(index)
|
||||
new_index[axis] = indices[index]
|
||||
out[index] = data[tuple(new_index)]
|
||||
return out
|
||||
@@ -0,0 +1,69 @@
|
||||
# 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.
|
||||
"""Numpy reference implementation for get_valid_counts."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_valid_counts_python(data, score_threshold=0, id_index=0, score_index=1):
|
||||
"""Numpy reference for get_valid_counts.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
3-D array with shape [batch_size, num_anchors, elem_length].
|
||||
|
||||
score_threshold : float
|
||||
Lower limit of score for valid bounding boxes.
|
||||
|
||||
id_index : int
|
||||
Index of the class categories, -1 to disable.
|
||||
|
||||
score_index : int
|
||||
Index of the scores/confidence of boxes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
valid_count : numpy.ndarray
|
||||
1-D array, shape [batch_size].
|
||||
|
||||
out_tensor : numpy.ndarray
|
||||
Rearranged data, shape [batch_size, num_anchors, elem_length].
|
||||
|
||||
out_indices : numpy.ndarray
|
||||
Indices mapping, shape [batch_size, num_anchors].
|
||||
"""
|
||||
batch_size, num_anchors, box_data_length = data.shape
|
||||
valid_count = np.zeros(batch_size, dtype="int32")
|
||||
out_tensor = np.full_like(data, -1.0)
|
||||
out_indices = np.full((batch_size, num_anchors), -1, dtype="int32")
|
||||
|
||||
for i in range(batch_size):
|
||||
cnt = 0
|
||||
for j in range(num_anchors):
|
||||
score = data[i, j, score_index]
|
||||
if id_index < 0:
|
||||
is_valid = score > score_threshold
|
||||
else:
|
||||
is_valid = score > score_threshold and data[i, j, id_index] >= 0
|
||||
if is_valid:
|
||||
out_tensor[i, cnt, :] = data[i, j, :]
|
||||
out_indices[i, cnt] = j
|
||||
cnt += 1
|
||||
valid_count[i] = cnt
|
||||
|
||||
return valid_count, out_tensor, out_indices
|
||||
@@ -0,0 +1,398 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""affine_grid and grid_sample operators in python"""
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def affine_grid_python(data, target_shape):
|
||||
yv, xv = np.meshgrid(np.arange(target_shape[0]), np.arange(target_shape[1]))
|
||||
yv = yv.T * 2 / (target_shape[0] - 1) - 1
|
||||
xv = xv.T * 2 / (target_shape[1] - 1) - 1
|
||||
ones = np.ones_like(xv)
|
||||
grid = np.stack([xv, yv, ones]).reshape(3, -1)
|
||||
return data.reshape(-1, 3).dot(grid).reshape(data.shape[0], 2, *target_shape)
|
||||
|
||||
|
||||
def grid_sample_2d(
|
||||
data: np.ndarray,
|
||||
grid: np.ndarray,
|
||||
method="bilinear",
|
||||
layout="NCHW",
|
||||
padding_mode="zeros",
|
||||
align_corners=True,
|
||||
):
|
||||
r"""grid_sample_2d for NCHW layout"""
|
||||
|
||||
assert method in ("bilinear", "nearest", "bicubic"), f"{method} is not supported"
|
||||
assert layout == "NCHW"
|
||||
assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
|
||||
assert len(data.shape) == len(grid.shape) == 4
|
||||
|
||||
batch, channel = data.shape[:2]
|
||||
in_height, in_width = data.shape[2:]
|
||||
out_height, out_width = grid.shape[2:]
|
||||
out_shape = [batch, channel, out_height, out_width]
|
||||
out = np.zeros(out_shape)
|
||||
|
||||
def _get_pixel(b, c, h, w):
|
||||
if 0 <= h <= in_height - 1 and 0 <= w <= in_width - 1:
|
||||
return data[b, c, h, w]
|
||||
return 0
|
||||
|
||||
def _unnormalize(h, w):
|
||||
if align_corners:
|
||||
new_h = (h + 1) * (in_height - 1) / 2
|
||||
new_w = (w + 1) * (in_width - 1) / 2
|
||||
else:
|
||||
new_h = -0.5 + (h + 1) * in_height / 2
|
||||
new_w = -0.5 + (w + 1) * in_width / 2
|
||||
return (new_h, new_w)
|
||||
|
||||
def _clip_coordinates(x, size):
|
||||
return min(max(x, 0), size - 1)
|
||||
|
||||
def _reflect_coordinates(i, size):
|
||||
def __refelection(i, size, corner_start):
|
||||
def __reflect(index, size, corner_start):
|
||||
index_align_corner = abs(corner_start - index)
|
||||
size_times = index_align_corner // size
|
||||
even = size_times % 2 == 0
|
||||
extra = index_align_corner - size_times * size
|
||||
return extra + corner_start if even else size - extra + corner_start
|
||||
|
||||
if corner_start <= i <= size + corner_start:
|
||||
new_i = i
|
||||
else:
|
||||
new_i = __reflect(i, size, corner_start)
|
||||
return new_i
|
||||
|
||||
if align_corners:
|
||||
x = __refelection(i, size - 1, 0)
|
||||
else:
|
||||
x = __refelection(i, size, -0.5)
|
||||
return x
|
||||
|
||||
def _compute_source_index(b, h, w):
|
||||
y = grid[b, 1, h, w]
|
||||
x = grid[b, 0, h, w]
|
||||
y, x = _unnormalize(y, x)
|
||||
|
||||
if padding_mode == "reflection":
|
||||
y = _reflect_coordinates(y, in_height)
|
||||
x = _reflect_coordinates(x, in_width)
|
||||
y = _clip_coordinates(y, in_height)
|
||||
x = _clip_coordinates(x, in_width)
|
||||
elif padding_mode == "border":
|
||||
y = _clip_coordinates(y, in_height)
|
||||
x = _clip_coordinates(x, in_width)
|
||||
|
||||
return (y, x)
|
||||
|
||||
def _nearest_sample():
|
||||
for _b in range(batch):
|
||||
for _c in range(channel):
|
||||
for _h in range(out_height):
|
||||
for _w in range(out_width):
|
||||
y, x = _compute_source_index(_b, _h, _w)
|
||||
# python round is not used here,
|
||||
# beacause it is done toward the even choice
|
||||
new_y = int(y + 0.5) if y > 0 else int(y - 0.5)
|
||||
new_x = int(x + 0.5) if x > 0 else int(x - 0.5)
|
||||
out[_b, _c, _h, _w] = _get_pixel(_b, _c, new_y, new_x)
|
||||
|
||||
def _bilinear_sample():
|
||||
for _b in range(batch):
|
||||
for _c in range(channel):
|
||||
for _h in range(out_height):
|
||||
for _w in range(out_width):
|
||||
y, x = _compute_source_index(_b, _h, _w)
|
||||
y0 = math.floor(y)
|
||||
x0 = math.floor(x)
|
||||
y1 = y0 + 1
|
||||
x1 = x0 + 1
|
||||
|
||||
out[_b, _c, _h, _w] = (
|
||||
_get_pixel(_b, _c, y0, x0) * (1.0 - (y - y0)) * (1.0 - (x - x0))
|
||||
+ _get_pixel(_b, _c, y0, x1) * (1.0 - (y - y0)) * (x - x0)
|
||||
+ _get_pixel(_b, _c, y1, x0) * (y - y0) * (1.0 - (x - x0))
|
||||
+ _get_pixel(_b, _c, y1, x1) * (y - y0) * (x - x0)
|
||||
)
|
||||
|
||||
def _bicubic_sample():
|
||||
A = -0.75
|
||||
|
||||
def cubic_weight_1(x_fraction):
|
||||
return ((A + 2) * x_fraction - (A + 3)) * x_fraction * x_fraction + 1
|
||||
|
||||
def cubic_weight_2(x_fraction):
|
||||
return ((A * x_fraction - 5 * A) * x_fraction + 8 * A) * x_fraction - 4 * A
|
||||
|
||||
def cubic_interp_1d(pixel_0, pixel_1, pixel_2, pixel_3, x_fraction):
|
||||
weights = [0] * 4
|
||||
weights[0] = cubic_weight_2(x_fraction + 1)
|
||||
weights[1] = cubic_weight_1(x_fraction)
|
||||
weights[2] = cubic_weight_1(1 - x_fraction)
|
||||
weights[3] = cubic_weight_2(2 - x_fraction)
|
||||
|
||||
return (
|
||||
pixel_0 * weights[0]
|
||||
+ pixel_1 * weights[1]
|
||||
+ pixel_2 * weights[2]
|
||||
+ pixel_3 * weights[3]
|
||||
)
|
||||
|
||||
def coefficients_along_x(x_floor, y_floor, x_fraction):
|
||||
coefficients = [0] * 4
|
||||
|
||||
for i in range(4):
|
||||
y_ = y_floor - 1 + i
|
||||
x_0 = x_floor - 1
|
||||
x_1 = x_floor + 0
|
||||
x_2 = x_floor + 1
|
||||
x_3 = x_floor + 2
|
||||
|
||||
if padding_mode == "border":
|
||||
y_ = _clip_coordinates(y_, in_height)
|
||||
x_0 = _clip_coordinates(x_0, in_width)
|
||||
x_1 = _clip_coordinates(x_1, in_width)
|
||||
x_2 = _clip_coordinates(x_2, in_width)
|
||||
x_3 = _clip_coordinates(x_3, in_width)
|
||||
|
||||
elif padding_mode == "reflection":
|
||||
y_ = _reflect_coordinates(y_, in_height)
|
||||
x_0 = _reflect_coordinates(x_0, in_width)
|
||||
x_1 = _reflect_coordinates(x_1, in_width)
|
||||
x_2 = _reflect_coordinates(x_2, in_width)
|
||||
x_3 = _reflect_coordinates(x_3, in_width)
|
||||
|
||||
y_ = int(_clip_coordinates(y_, in_height))
|
||||
x_0 = int(_clip_coordinates(x_0, in_width))
|
||||
x_1 = int(_clip_coordinates(x_1, in_width))
|
||||
x_2 = int(_clip_coordinates(x_2, in_width))
|
||||
x_3 = int(_clip_coordinates(x_3, in_width))
|
||||
|
||||
coefficients[i] = cubic_interp_1d(
|
||||
_get_pixel(_b, _c, y_, x_0),
|
||||
_get_pixel(_b, _c, y_, x_1),
|
||||
_get_pixel(_b, _c, y_, x_2),
|
||||
_get_pixel(_b, _c, y_, x_3),
|
||||
x_fraction,
|
||||
)
|
||||
return coefficients
|
||||
|
||||
for _b in range(batch):
|
||||
for _c in range(channel):
|
||||
for _h in range(out_height):
|
||||
for _w in range(out_width):
|
||||
y = grid[_b, 1, _h, _w]
|
||||
x = grid[_b, 0, _h, _w]
|
||||
y, x = _unnormalize(y, x)
|
||||
y_floor = math.floor(y)
|
||||
x_floor = math.floor(x)
|
||||
y_fraction = y - y_floor
|
||||
x_fraction = x - x_floor
|
||||
|
||||
coefficients = coefficients_along_x(x_floor, y_floor, x_fraction)
|
||||
|
||||
out[_b, _c, _h, _w] = cubic_interp_1d(
|
||||
coefficients[0],
|
||||
coefficients[1],
|
||||
coefficients[2],
|
||||
coefficients[3],
|
||||
y_fraction,
|
||||
)
|
||||
|
||||
if method == "bilinear":
|
||||
_bilinear_sample()
|
||||
elif method == "nearest":
|
||||
_nearest_sample()
|
||||
else: # mode == "bicubic":
|
||||
_bicubic_sample()
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def grid_sample_3d(
|
||||
data: np.ndarray,
|
||||
grid: np.ndarray,
|
||||
method="bilinear",
|
||||
layout="NCDHW",
|
||||
padding_mode="zeros",
|
||||
align_corners=True,
|
||||
):
|
||||
r"""grid_sample_3d for NCDHW layout"""
|
||||
|
||||
assert method in ("bilinear", "nearest"), f"{method} is not supported"
|
||||
assert layout == "NCDHW"
|
||||
assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
|
||||
assert len(data.shape) == len(grid.shape) == 5
|
||||
|
||||
batch, channel = data.shape[:2]
|
||||
in_depth, in_height, in_width = data.shape[2:]
|
||||
out_depth, out_height, out_width = grid.shape[2:]
|
||||
out_shape = [batch, channel, out_depth, out_height, out_width]
|
||||
out = np.zeros(out_shape)
|
||||
|
||||
def _get_pixel(b, c, d, h, w):
|
||||
if 0 <= d <= in_depth - 1 and 0 <= h <= in_height - 1 and 0 <= w <= in_width - 1:
|
||||
return data[b, c, d, h, w]
|
||||
return 0
|
||||
|
||||
def _unnormalize(d, h, w):
|
||||
if align_corners:
|
||||
new_d = (d + 1) * (in_depth - 1) / 2
|
||||
new_h = (h + 1) * (in_height - 1) / 2
|
||||
new_w = (w + 1) * (in_width - 1) / 2
|
||||
else:
|
||||
new_d = -0.5 + (d + 1) * in_depth / 2
|
||||
new_h = -0.5 + (h + 1) * in_height / 2
|
||||
new_w = -0.5 + (w + 1) * in_width / 2
|
||||
return (new_d, new_h, new_w)
|
||||
|
||||
def _clip_coordinates(x, size):
|
||||
return min(max(x, 0), size - 1)
|
||||
|
||||
def _reflect_coordinates(i, size):
|
||||
def __refelection(i, size, corner_start):
|
||||
def __reflect(index, size, corner_start):
|
||||
index_align_corner = abs(corner_start - index)
|
||||
size_times = index_align_corner // size
|
||||
even = size_times % 2 == 0
|
||||
extra = index_align_corner - size_times * size
|
||||
return extra + corner_start if even else size - extra + corner_start
|
||||
|
||||
if corner_start <= i <= size + corner_start:
|
||||
new_i = i
|
||||
else:
|
||||
new_i = __reflect(i, size, corner_start)
|
||||
return new_i
|
||||
|
||||
if align_corners:
|
||||
x = __refelection(i, size - 1, 0)
|
||||
else:
|
||||
x = __refelection(i, size, -0.5)
|
||||
return x
|
||||
|
||||
def _compute_source_index(b, d, h, w):
|
||||
z = grid[b, 2, d, h, w]
|
||||
y = grid[b, 1, d, h, w]
|
||||
x = grid[b, 0, d, h, w]
|
||||
z, y, x = _unnormalize(z, y, x)
|
||||
|
||||
if padding_mode == "reflection":
|
||||
z = _reflect_coordinates(z, in_depth)
|
||||
y = _reflect_coordinates(y, in_height)
|
||||
x = _reflect_coordinates(x, in_width)
|
||||
z = _clip_coordinates(z, in_depth)
|
||||
y = _clip_coordinates(y, in_height)
|
||||
x = _clip_coordinates(x, in_width)
|
||||
elif padding_mode == "border":
|
||||
z = _clip_coordinates(z, in_depth)
|
||||
y = _clip_coordinates(y, in_height)
|
||||
x = _clip_coordinates(x, in_width)
|
||||
return (z, y, x)
|
||||
|
||||
def _nearest_sample():
|
||||
for _b in range(batch):
|
||||
for _c in range(channel):
|
||||
for _d in range(out_depth):
|
||||
for _h in range(out_height):
|
||||
for _w in range(out_width):
|
||||
z, y, x = _compute_source_index(_b, _d, _h, _w)
|
||||
# python round is not used here,
|
||||
# beacause it is done toward the even choice
|
||||
new_z = int(z + 0.5) if z > 0 else int(z - 0.5)
|
||||
new_y = int(y + 0.5) if y > 0 else int(y - 0.5)
|
||||
new_x = int(x + 0.5) if x > 0 else int(x - 0.5)
|
||||
out[_b, _c, _d, _h, _w] = _get_pixel(_b, _c, new_z, new_y, new_x)
|
||||
|
||||
def _triilinear_sample():
|
||||
for _b in range(batch):
|
||||
for _c in range(channel):
|
||||
for _d in range(out_depth):
|
||||
for _h in range(out_height):
|
||||
for _w in range(out_width):
|
||||
z, y, x = _compute_source_index(_b, _d, _h, _w)
|
||||
z0 = math.floor(z)
|
||||
y0 = math.floor(y)
|
||||
x0 = math.floor(x)
|
||||
z1 = z0 + 1
|
||||
y1 = y0 + 1
|
||||
x1 = x0 + 1
|
||||
|
||||
out[_b, _c, _d, _h, _w] = (
|
||||
_get_pixel(_b, _c, z0, y0, x0)
|
||||
* (1 - (x - x0))
|
||||
* (1 - (y - y0))
|
||||
* (1 - (z - z0))
|
||||
+ _get_pixel(_b, _c, z0, y0, x1)
|
||||
* (x - x0)
|
||||
* (1 - (y - y0))
|
||||
* (1 - (z - z0))
|
||||
+ _get_pixel(_b, _c, z1, y1, x0)
|
||||
* (1 - (x - x0))
|
||||
* (y - y0)
|
||||
* (z - z0)
|
||||
+ _get_pixel(_b, _c, z1, y1, x1) * (x - x0) * (y - y0) * (z - z0)
|
||||
+ _get_pixel(_b, _c, z0, y1, x0)
|
||||
* (1 - (x - x0))
|
||||
* (y - y0)
|
||||
* (1 - (z - z0))
|
||||
+ _get_pixel(_b, _c, z1, y0, x1)
|
||||
* (x - x0)
|
||||
* (1 - (y - y0))
|
||||
* (z - z0)
|
||||
+ _get_pixel(_b, _c, z1, y0, x0)
|
||||
* (1 - (x - x0))
|
||||
* (1 - (y - y0))
|
||||
* (z - z0)
|
||||
+ _get_pixel(_b, _c, z0, y1, x1)
|
||||
* (x - x0)
|
||||
* (y - y0)
|
||||
* (1 - (z - z0))
|
||||
)
|
||||
|
||||
if method == "bilinear":
|
||||
_triilinear_sample()
|
||||
else: # method == "nearest":
|
||||
_nearest_sample()
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def grid_sample_python(
|
||||
data: np.ndarray,
|
||||
grid: np.ndarray,
|
||||
method="bilinear",
|
||||
layout="NCHW",
|
||||
padding_mode="zeros",
|
||||
align_corners=True,
|
||||
):
|
||||
r"""grid_sample_3d for NCDHW layout or grid_sample_2d for NCHW layout"""
|
||||
|
||||
if len(data.shape) == 4:
|
||||
grid_sample = grid_sample_2d
|
||||
elif len(data.shape) == 5:
|
||||
grid_sample = grid_sample_3d
|
||||
else:
|
||||
raise ValueError("invalid shape")
|
||||
|
||||
return grid_sample(data, grid, method, layout, padding_mode, align_corners)
|
||||
@@ -0,0 +1,84 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Group normalization in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def group_norm_python(data, gamma, beta, num_groups, channel_axis, axes, epsilon=1e-5):
|
||||
"""Group normalization operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
gamma: tvm.te.Tensor
|
||||
1-D with shape (r_0) where r_0 == d_{channel_axis}
|
||||
|
||||
beta: tvm.te.Tensor
|
||||
Optional, 1-D with shape (r_0) where r_0 == d_{channel_axis}
|
||||
|
||||
num_groups : int
|
||||
The number of groups
|
||||
|
||||
channel_axis : int
|
||||
The channel axis
|
||||
|
||||
axes : list of int
|
||||
Axis over the normalization applied, excluding the channel axis
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
old_shape = data.shape
|
||||
old_dtype = data.dtype
|
||||
new_shape = list(old_shape)
|
||||
new_shape[channel_axis] = data.shape[channel_axis] // num_groups
|
||||
new_shape.insert(channel_axis, num_groups)
|
||||
data = np.reshape(data, new_shape).astype("float32")
|
||||
new_axes = [channel_axis + 1]
|
||||
for axis in axes:
|
||||
if axis < channel_axis:
|
||||
new_axes.append(axis)
|
||||
else:
|
||||
new_axes.append(axis + 1)
|
||||
mean = np.mean(data, axis=tuple(new_axes), keepdims=True)
|
||||
var = np.var(data, axis=tuple(new_axes), keepdims=True)
|
||||
data = (data - mean) / np.sqrt(var + epsilon)
|
||||
data = np.reshape(data, old_shape).astype(old_dtype)
|
||||
|
||||
gamma_broadcast_shape = [1 for _ in range(len(old_shape))]
|
||||
gamma_broadcast_shape[channel_axis] = gamma.shape[0]
|
||||
gamma = np.reshape(gamma, gamma_broadcast_shape)
|
||||
|
||||
beta_broadcast_shape = [1 for _ in range(len(old_shape))]
|
||||
beta_broadcast_shape[channel_axis] = beta.shape[0]
|
||||
if beta is not None:
|
||||
beta = np.reshape(beta, beta_broadcast_shape)
|
||||
|
||||
data *= gamma
|
||||
if beta is not None:
|
||||
data += beta
|
||||
|
||||
return data
|
||||
@@ -0,0 +1,54 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Instance normalization in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def instance_norm_python(data, gamma, beta, axis, epsilon=1e-5):
|
||||
"""Instance normalization operator in Python.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
gamma: numpy.ndarray
|
||||
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
beta: numpy.ndarray
|
||||
Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
axis : int or tuple of ints
|
||||
Axis over the normalization applied
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : np.ndarray
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
mean = np.mean(data, axis, keepdims=True)
|
||||
var = np.var(data, axis, keepdims=True)
|
||||
result = (data - mean) / np.sqrt(var + epsilon)
|
||||
result *= gamma
|
||||
if beta is not None:
|
||||
result += beta
|
||||
return result
|
||||
@@ -0,0 +1,45 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""L2 normalize in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def l2_normalize_python(a_np, eps, axis=None):
|
||||
"""L2 normalize operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
eps : float
|
||||
epsilon constant value
|
||||
axis : list of int
|
||||
axis over the normalization applied
|
||||
|
||||
Returns
|
||||
-------
|
||||
l2_normalize_out : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
dot_value = np.power(a_np, 2.0)
|
||||
sqr_sum = np.sum(dot_value, axis, keepdims=True)
|
||||
sqrt_sum = np.sqrt(np.maximum(np.broadcast_to(sqr_sum, a_np.shape), eps))
|
||||
l2_normalize_out = np.divide(a_np, sqrt_sum)
|
||||
return l2_normalize_out
|
||||
@@ -0,0 +1,57 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Layer normalization in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def layer_norm_python(data, gamma, beta, axis, epsilon=1e-5):
|
||||
"""Layer normalization operator in Python.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
gamma: numpy.ndarray
|
||||
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
beta: numpy.ndarray
|
||||
Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
axis : int or tuple of ints
|
||||
Axis over the normalization applied
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : np.ndarray
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
old_dtype = data.dtype
|
||||
data = data.astype("float32")
|
||||
mean = np.mean(data, axis, keepdims=True)
|
||||
var = np.var(data, axis, keepdims=True)
|
||||
result = (data - mean) / np.sqrt(var + epsilon)
|
||||
result = result.astype(old_dtype)
|
||||
result *= gamma
|
||||
if beta is not None:
|
||||
result += beta
|
||||
return result
|
||||
@@ -0,0 +1,75 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
# ruff: noqa: E741
|
||||
"""LRN in python"""
|
||||
|
||||
from itertools import product
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def lrn_python(a_np, size, axis, bias, alpha, beta):
|
||||
"""Local response normalization operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
size : int
|
||||
normalization window size
|
||||
|
||||
axis : int
|
||||
input data layout channel axis
|
||||
|
||||
bias : float
|
||||
offset to avoid dividing by 0. constant value
|
||||
|
||||
alpha : float
|
||||
constant value
|
||||
|
||||
beta : float
|
||||
exponent constant value
|
||||
|
||||
Returns
|
||||
-------
|
||||
lrn_out : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
radius = size // 2
|
||||
sqr_sum = np.zeros(shape=a_np.shape).astype(a_np.dtype)
|
||||
for i, j, k, l in product(*[range(_axis) for _axis in a_np.shape]):
|
||||
axis_size = a_np.shape[axis]
|
||||
if axis == 1:
|
||||
# NCHW layout
|
||||
sum_start = j - radius if j - radius >= 0 else 0
|
||||
sum_end = j + radius + 1 if j + radius + 1 < axis_size else axis_size
|
||||
sqr_sum[i, j, k, l] = sum(
|
||||
a_np[i, sum_start:sum_end, k, l] * a_np[i, sum_start:sum_end, k, l]
|
||||
)
|
||||
elif axis == 3:
|
||||
# NHWC layout
|
||||
sum_start = l - radius if l - radius >= 0 else 0
|
||||
sum_end = l + radius + 1 if l + radius + 1 < axis_size else axis_size
|
||||
sqr_sum[i, j, k, l] = sum(
|
||||
a_np[i, j, k, sum_start:sum_end] * a_np[i, j, k, sum_start:sum_end]
|
||||
)
|
||||
|
||||
sqr_sum_up = np.power((bias + (alpha * sqr_sum / size)), beta)
|
||||
lrn_out = np.divide(a_np, sqr_sum_up)
|
||||
return lrn_out
|
||||
@@ -0,0 +1,136 @@
|
||||
# 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
|
||||
# ruff: noqa: E741
|
||||
"""LSTM reference implementation using numpy."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def lstm_python(
|
||||
Xs: np.array,
|
||||
Wi: np.array,
|
||||
Wh: np.array,
|
||||
Bi: np.array = None,
|
||||
Bh: np.array = None,
|
||||
h_init: np.array = None,
|
||||
c_init: np.array = None,
|
||||
proj: np.array = None,
|
||||
p_i: np.array = None,
|
||||
p_f: np.array = None,
|
||||
p_o: np.array = None,
|
||||
f_act: str = "sigmoid",
|
||||
g_act: str = "tanh",
|
||||
h_act: str = "tanh",
|
||||
reverse: bool = False,
|
||||
weight_layout: str = "IFGO",
|
||||
):
|
||||
"""LSTM reference implementation using numpy
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Xs : np.array
|
||||
(seq_length, batch_size, in_dim)
|
||||
Wi : np.array
|
||||
(4 * hidden_dim, in_dim)
|
||||
Wh : np.array
|
||||
(4 * hidden_dim, out_dim) where out_dim = proj_dim if proj_dim > 0, else hidden_dim
|
||||
Bi : np.array, optional
|
||||
(4 * hidden_dim,), by default None
|
||||
Bh : np.array, optional
|
||||
(4 * hidden_dim,), by default None
|
||||
h_init : np.array, optional
|
||||
(batch_size, out_dim), by default None
|
||||
c_init : np.array, optional
|
||||
(batch_size, hidden_dim), by default None
|
||||
proj : np.array, optional
|
||||
(proj_dim, hidden_dim), by default None
|
||||
p_i, p_f, p_o: np.array, optional
|
||||
(batch_size, hidden_dim), by default None
|
||||
f_act, g_act, h_act: str, optional
|
||||
activations, by default "sigmoid", "tanh", "tanh"
|
||||
reverse : bool, optional
|
||||
process Xs in reverse, by default False
|
||||
weight_layout : str, optional
|
||||
Packed layout for weights and biases, by default "IFGO"
|
||||
"""
|
||||
i_gate_idx = weight_layout.find("I")
|
||||
f_gate_idx = weight_layout.find("F")
|
||||
g_gate_idx = weight_layout.find("G")
|
||||
o_gate_idx = weight_layout.find("O")
|
||||
|
||||
str2act = {"sigmoid": lambda x: 1 / (1 + np.exp(-x)), "tanh": np.tanh}
|
||||
|
||||
f_act = str2act[f_act]
|
||||
g_act = str2act[g_act]
|
||||
h_act = str2act[h_act]
|
||||
|
||||
S, B, F = Xs.shape
|
||||
H = Wi.shape[0] // 4
|
||||
O = Wh.shape[1]
|
||||
|
||||
# make life a bit easier
|
||||
Wi = np.reshape(Wi, (4, H, F))
|
||||
Wh = np.reshape(Wh, (4, H, O))
|
||||
if Bi is not None:
|
||||
Bi = np.reshape(Bi, (4, H))
|
||||
if Bh is not None:
|
||||
Bh = np.reshape(Bh, (4, H))
|
||||
|
||||
h0 = h_init if h_init is not None else np.zeros((B, O), "float32")
|
||||
c0 = c_init if c_init is not None else np.zeros((B, H), "float32")
|
||||
|
||||
hs = [h0]
|
||||
cs = [c0]
|
||||
|
||||
for t in range(S):
|
||||
x = Xs[S - t - 1 if reverse else t]
|
||||
xh = [np.matmul(x, Wi[g].T) for g in range(4)]
|
||||
if Bi is not None:
|
||||
xh = [xh[g] + Bi[g] for g in range(4)]
|
||||
|
||||
hh = [np.matmul(hs[t], Wh[g].T) for g in range(4)]
|
||||
if Bh is not None:
|
||||
hh = [hh[g] + Bh[g] for g in range(4)]
|
||||
|
||||
sums = [xh[g] + hh[g] for g in range(4)]
|
||||
|
||||
if p_i is not None and p_f is not None:
|
||||
i_gate = f_act(sums[i_gate_idx] + p_i * cs[t])
|
||||
f_gate = f_act(sums[f_gate_idx] + p_f * cs[t])
|
||||
else:
|
||||
i_gate = f_act(sums[i_gate_idx])
|
||||
f_gate = f_act(sums[f_gate_idx])
|
||||
|
||||
g_gate = g_act(sums[g_gate_idx])
|
||||
|
||||
next_c = f_gate * cs[t] + i_gate * g_gate
|
||||
|
||||
if p_o is not None:
|
||||
o_gate = f_act(sums[o_gate_idx] + p_o * next_c)
|
||||
else:
|
||||
o_gate = f_act(sums[o_gate_idx])
|
||||
|
||||
next_h = o_gate * h_act(next_c)
|
||||
|
||||
if proj is not None:
|
||||
next_h = np.matmul(next_h, proj.T)
|
||||
|
||||
hs.append(next_h)
|
||||
cs.append(next_c)
|
||||
|
||||
return np.stack(hs[1:], axis=0), np.stack(cs[1:], axis=0)
|
||||
@@ -0,0 +1,85 @@
|
||||
# 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
|
||||
"""MatrixSetDiag in Python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def matrix_set_diag(input_np, diagonal, k=0, align="RIGHT_LEFT"):
|
||||
"""matrix_set_diag operator implemented in numpy.
|
||||
|
||||
Returns a numpy array with the diagonals of input array
|
||||
replaced with the provided diagonal values.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_np : numpy.ndarray
|
||||
Input Array.
|
||||
Shape = [D1, D2, D3, ... , Dn-1 , Dn]
|
||||
|
||||
diagonal : numpy.ndarray
|
||||
Values to be filled in the diagonal.
|
||||
|
||||
k : int or tuple of int
|
||||
Diagonal Offsets.
|
||||
|
||||
align : string
|
||||
Some diagonals are shorter than max_diag_len and need to be padded.
|
||||
Possible Vales:
|
||||
["RIGHT_LEFT" (default), "LEFT_RIGHT", "LEFT_LEFT", "RIGHT_RIGHT"]
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : numpy.ndarray
|
||||
New Array with given diagonal values.
|
||||
Shape = [D1, D2, D3, ... , Dn-1 , Dn]
|
||||
"""
|
||||
out = np.array(input_np, copy=True)
|
||||
|
||||
cols = input_np.shape[-1]
|
||||
rows = input_np.shape[-2]
|
||||
|
||||
onlyOneDiagonal = True
|
||||
if isinstance(k, tuple | list):
|
||||
if len(k) < 2 or k[0] == k[1]:
|
||||
k = k[0]
|
||||
else:
|
||||
onlyOneDiagonal = False
|
||||
|
||||
if onlyOneDiagonal:
|
||||
for i in range(diagonal.shape[-1]):
|
||||
if k >= 0:
|
||||
out[..., i, i + k] = diagonal[..., i]
|
||||
else:
|
||||
out[..., i - k, i] = diagonal[..., i]
|
||||
else:
|
||||
for ki in range(k[0], k[1] + 1):
|
||||
diag_len = min(cols - max(ki, 0), rows + min(ki, 0))
|
||||
offset = 0
|
||||
if ki >= 0:
|
||||
if align[:5] == "RIGHT":
|
||||
offset = diagonal.shape[-1] - diag_len
|
||||
else:
|
||||
if align[-5:] == "RIGHT":
|
||||
offset = diagonal.shape[-1] - diag_len
|
||||
for i in range(diag_len):
|
||||
if ki >= 0:
|
||||
out[..., i, i + ki] = diagonal[..., k[1] - ki, i + offset]
|
||||
else:
|
||||
out[..., i - ki, i] = diagonal[..., k[1] - ki, i + offset]
|
||||
return out
|
||||
@@ -0,0 +1,74 @@
|
||||
# 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
|
||||
# ruff: noqa: RUF005
|
||||
"""NLLLoss in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def nll_loss(predictions, targets, weights, reduction="mean", ignore_index=-100):
|
||||
"""nll_loss operator implemented in numpy.
|
||||
|
||||
output{n, i_1, i_2, ..., i_k} = -p * w
|
||||
where t = target{n, i_1, i_2, ..., i_k}
|
||||
p = predictions{n, t, i_1, i_2, i_k}
|
||||
w = weights{n, i_1, i_2, ..., i_k} if t != ignore_index else 0
|
||||
|
||||
result = reduction(output)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
predictions : numpy.ndarray
|
||||
(k+2)-D with shape (N, C, d_1, d_2, ..., d_k),
|
||||
where C is the number of target classes
|
||||
|
||||
targets : numpy.ndarray
|
||||
(k+1)-D with shape (N, d_1, d_2, ..., d_k)
|
||||
The target value of the input.
|
||||
|
||||
weights : numpy.ndarray
|
||||
1-D with shape (C,)
|
||||
The weight of each target value.
|
||||
|
||||
reduction : string
|
||||
The reduction method to apply to output.
|
||||
Can be "mean", "sum" or "none".
|
||||
|
||||
ignore_index : int
|
||||
The target value to ignore.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : numpy.ndarray
|
||||
a scalar if the reduction type is "mean" or "sum",
|
||||
otherwise the same shape as `target`.
|
||||
"""
|
||||
res = np.zeros(targets.shape)
|
||||
weight_sum = 0.0
|
||||
for index in np.ndindex(targets.shape):
|
||||
class_id = targets[index]
|
||||
if class_id != ignore_index:
|
||||
index_list = list(index)
|
||||
pred_index = tuple(index_list[:1] + [class_id] + index_list[1:])
|
||||
res[index] = -predictions[pred_index] * weights[class_id]
|
||||
weight_sum += weights[class_id]
|
||||
if reduction == "mean":
|
||||
return np.sum(res) / weight_sum
|
||||
if reduction == "sum":
|
||||
return np.sum(res)
|
||||
return res
|
||||
@@ -0,0 +1,219 @@
|
||||
# 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.
|
||||
"""Numpy reference implementation for classic non_max_suppression."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _iou(box_a, box_b, coord_start):
|
||||
"""Compute IoU between two boxes."""
|
||||
a = box_a[coord_start : coord_start + 4]
|
||||
b = box_b[coord_start : coord_start + 4]
|
||||
|
||||
a_l, a_t, a_r, a_b = min(a[0], a[2]), min(a[1], a[3]), max(a[0], a[2]), max(a[1], a[3])
|
||||
b_l, b_t, b_r, b_b = min(b[0], b[2]), min(b[1], b[3]), max(b[0], b[2]), max(b[1], b[3])
|
||||
|
||||
w = max(0.0, min(a_r, b_r) - max(a_l, b_l))
|
||||
h = max(0.0, min(a_b, b_b) - max(a_t, b_t))
|
||||
area = w * h
|
||||
u = (a_r - a_l) * (a_b - a_t) + (b_r - b_l) * (b_b - b_t) - area
|
||||
return 0.0 if u <= 0 else area / u
|
||||
|
||||
|
||||
def non_max_suppression_python(
|
||||
data,
|
||||
valid_count,
|
||||
indices,
|
||||
max_output_size=-1,
|
||||
iou_threshold=0.5,
|
||||
force_suppress=False,
|
||||
top_k=-1,
|
||||
coord_start=2,
|
||||
score_index=1,
|
||||
id_index=0,
|
||||
return_indices=True,
|
||||
invalid_to_bottom=False,
|
||||
soft_nms_sigma=0.0,
|
||||
score_threshold=0.0,
|
||||
):
|
||||
"""Numpy reference for classic non_max_suppression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
3-D array, shape [batch_size, num_anchors, elem_length].
|
||||
|
||||
valid_count : numpy.ndarray
|
||||
1-D array, shape [batch_size].
|
||||
|
||||
indices : numpy.ndarray
|
||||
2-D array, shape [batch_size, num_anchors].
|
||||
|
||||
Returns
|
||||
-------
|
||||
If return_indices is True and soft_nms_sigma == 0.0: (box_indices, valid_box_count)
|
||||
If return_indices is True and soft_nms_sigma > 0.0:
|
||||
(out_data, box_indices, valid_box_count)
|
||||
Otherwise: modified data tensor
|
||||
"""
|
||||
batch_size, num_anchors, _ = data.shape
|
||||
out_data = np.full_like(data, -1.0)
|
||||
out_box_indices = np.full((batch_size, num_anchors), -1, dtype="int32")
|
||||
compacted = np.full((batch_size, num_anchors), -1, dtype="int32")
|
||||
valid_box_count = np.zeros((batch_size, 1), dtype="int32")
|
||||
|
||||
is_soft_nms = soft_nms_sigma > 0.0
|
||||
thresh = score_threshold if is_soft_nms else 0.0
|
||||
soft_nms_scale = -0.5 / soft_nms_sigma if is_soft_nms else 0.0
|
||||
|
||||
for i in range(batch_size):
|
||||
nkeep = int(valid_count[i])
|
||||
if 0 < top_k < nkeep:
|
||||
nkeep = top_k
|
||||
|
||||
# Sort by score descending
|
||||
scores = data[i, :nkeep, score_index].copy()
|
||||
sorted_idx = np.argsort(-scores)
|
||||
|
||||
# Copy sorted boxes
|
||||
for j in range(nkeep):
|
||||
src = sorted_idx[j]
|
||||
out_data[i, j, :] = data[i, src, :]
|
||||
out_box_indices[i, j] = src
|
||||
|
||||
if is_soft_nms:
|
||||
num_selected = 0
|
||||
while num_selected < nkeep and (max_output_size < 0 or num_selected < max_output_size):
|
||||
best_idx = -1
|
||||
best_score = thresh
|
||||
for j in range(num_selected, nkeep):
|
||||
if out_box_indices[i, j] >= 0 and out_data[i, j, score_index] > best_score:
|
||||
best_idx = j
|
||||
best_score = out_data[i, j, score_index]
|
||||
|
||||
if best_idx < 0:
|
||||
break
|
||||
|
||||
if best_idx != num_selected:
|
||||
out_data[i, [num_selected, best_idx], :] = out_data[
|
||||
i, [best_idx, num_selected], :
|
||||
]
|
||||
out_box_indices[i, [num_selected, best_idx]] = out_box_indices[
|
||||
i, [best_idx, num_selected]
|
||||
]
|
||||
|
||||
selected_idx = num_selected
|
||||
for j in range(selected_idx + 1, nkeep):
|
||||
if out_box_indices[i, j] < 0 or out_data[i, j, score_index] <= thresh:
|
||||
continue
|
||||
|
||||
do_suppress = False
|
||||
if force_suppress:
|
||||
do_suppress = True
|
||||
elif id_index >= 0:
|
||||
do_suppress = (
|
||||
out_data[i, selected_idx, id_index] == out_data[i, j, id_index]
|
||||
)
|
||||
else:
|
||||
do_suppress = True
|
||||
|
||||
if not do_suppress:
|
||||
continue
|
||||
|
||||
iou = _iou(out_data[i, selected_idx], out_data[i, j], coord_start)
|
||||
if iou >= iou_threshold:
|
||||
out_box_indices[i, j] = -1
|
||||
else:
|
||||
out_data[i, j, score_index] *= np.exp(soft_nms_scale * (iou**2))
|
||||
if out_data[i, j, score_index] <= thresh:
|
||||
out_box_indices[i, j] = -1
|
||||
|
||||
num_selected += 1
|
||||
|
||||
valid_box_count[i, 0] = num_selected
|
||||
if return_indices:
|
||||
for j in range(num_selected):
|
||||
orig_idx = out_box_indices[i, j]
|
||||
compacted[i, j] = int(indices[i, orig_idx])
|
||||
out_box_indices[i, j] = compacted[i, j]
|
||||
for j in range(num_selected, num_anchors):
|
||||
out_data[i, j, :] = -1.0
|
||||
out_box_indices[i, j] = -1
|
||||
else:
|
||||
out_data[i, num_selected:, :] = -1.0
|
||||
continue
|
||||
|
||||
# Greedy NMS
|
||||
num_valid = 0
|
||||
for j in range(nkeep):
|
||||
if out_data[i, j, score_index] <= thresh:
|
||||
out_data[i, j, :] = -1.0
|
||||
out_box_indices[i, j] = -1
|
||||
continue
|
||||
if 0 < max_output_size <= num_valid:
|
||||
out_data[i, j, :] = -1.0
|
||||
out_box_indices[i, j] = -1
|
||||
continue
|
||||
|
||||
num_valid += 1
|
||||
|
||||
# Suppress overlapping boxes
|
||||
for k in range(j + 1, nkeep):
|
||||
if out_data[i, k, score_index] <= thresh:
|
||||
continue
|
||||
|
||||
do_suppress = False
|
||||
if force_suppress:
|
||||
do_suppress = True
|
||||
elif id_index >= 0:
|
||||
do_suppress = out_data[i, j, id_index] == out_data[i, k, id_index]
|
||||
else:
|
||||
do_suppress = True
|
||||
|
||||
if do_suppress:
|
||||
iou = _iou(out_data[i, j], out_data[i, k], coord_start)
|
||||
if iou >= iou_threshold:
|
||||
out_data[i, k, score_index] = -1.0
|
||||
out_box_indices[i, k] = -1
|
||||
|
||||
if return_indices:
|
||||
# Compact valid indices to top and remap to original
|
||||
cnt = 0
|
||||
for j in range(num_anchors):
|
||||
if out_box_indices[i, j] >= 0:
|
||||
orig_idx = out_box_indices[i, j]
|
||||
compacted[i, cnt] = int(indices[i, orig_idx])
|
||||
cnt += 1
|
||||
valid_box_count[i, 0] = cnt
|
||||
|
||||
if return_indices:
|
||||
if is_soft_nms:
|
||||
return [out_data, compacted, valid_box_count]
|
||||
return [compacted, valid_box_count]
|
||||
|
||||
if invalid_to_bottom:
|
||||
# Rearrange valid boxes to top
|
||||
result = np.full_like(data, -1.0)
|
||||
for i in range(batch_size):
|
||||
cnt = 0
|
||||
for j in range(num_anchors):
|
||||
if out_data[i, j, score_index] >= 0:
|
||||
result[i, cnt, :] = out_data[i, j, :]
|
||||
cnt += 1
|
||||
return result
|
||||
|
||||
return out_data
|
||||
@@ -0,0 +1,81 @@
|
||||
# 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
|
||||
"""OneHot in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def one_hot(indices, on_value, off_value, depth, axis, dtype):
|
||||
"""one_hot operator implemented in numpy.
|
||||
|
||||
Returns a one-hot tensor where the locations repsented by indices take value on_value,
|
||||
other locations take value off_value.
|
||||
Final dimension is <indices outer dimensions> x depth x <indices inner dimensions>.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices : numpy.ndarray
|
||||
Locations to set to on_value.
|
||||
|
||||
on_value : int/float
|
||||
Value to fill at indices.
|
||||
|
||||
off_value : int/float
|
||||
Value to fill at all other positions besides indices.
|
||||
|
||||
depth : int
|
||||
Depth of the one-hot dimension.
|
||||
|
||||
axis : int
|
||||
Axis to fill.
|
||||
|
||||
dtype : str
|
||||
Data type of the output tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : tvm.te.Tensor
|
||||
The one-hot tensor.
|
||||
"""
|
||||
oshape = []
|
||||
true_axis = len(indices.shape) if axis == -1 else axis
|
||||
ndim = len(indices.shape) + 1
|
||||
indices_index = 0
|
||||
for i in range(0, ndim):
|
||||
if i == true_axis:
|
||||
oshape.append(depth)
|
||||
else:
|
||||
oshape.append(indices.shape[indices_index])
|
||||
indices_index += 1
|
||||
|
||||
out = np.empty(oshape)
|
||||
output_indices = list(np.ndindex(out.shape))
|
||||
for output_index in output_indices:
|
||||
indices_indices = []
|
||||
for i, out_idx in enumerate(output_index):
|
||||
if i == true_axis:
|
||||
continue
|
||||
indices_indices.append(out_idx)
|
||||
|
||||
index = output_index[true_axis]
|
||||
if indices[tuple(indices_indices)] == index:
|
||||
out[output_index] = on_value
|
||||
else:
|
||||
out[output_index] = off_value
|
||||
|
||||
return out.astype(dtype)
|
||||
@@ -0,0 +1,71 @@
|
||||
# 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
|
||||
@@ -0,0 +1,209 @@
|
||||
# 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: RUF005
|
||||
"""Ground truth max and average pooling operators in python."""
|
||||
|
||||
import itertools
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm
|
||||
|
||||
|
||||
def _get_supported_layout(dims: int):
|
||||
"""
|
||||
Returns layout that is supported by poolnd_python based on number of
|
||||
dimensions of input tensor
|
||||
"""
|
||||
assert dims in [3, 4, 5], f"{dims}-dimensional tensor is not supported"
|
||||
if dims == 3:
|
||||
return "NCW"
|
||||
if dims == 4:
|
||||
return "NCHW"
|
||||
# dims == 5
|
||||
return "NCDHW"
|
||||
|
||||
|
||||
def _convert_to_layout(input_tensor: np.ndarray, layout: str) -> np.ndarray:
|
||||
"""
|
||||
Converts back to original layout after the algorithm is finished
|
||||
"""
|
||||
supported_layout = _get_supported_layout(input_tensor.ndim)
|
||||
if layout is not None and supported_layout != layout:
|
||||
# Generate transpose list
|
||||
transpose_list = []
|
||||
for d in layout:
|
||||
transpose_list.append(supported_layout.index(d))
|
||||
return input_tensor.transpose(transpose_list)
|
||||
return input_tensor
|
||||
|
||||
|
||||
def _convert_from_layout(input_tensor: np.ndarray, layout: str) -> np.ndarray:
|
||||
"""
|
||||
Converts tensor to one of suppored layouts
|
||||
"""
|
||||
supported_layout = _get_supported_layout(input_tensor.ndim)
|
||||
if layout is not None and supported_layout != layout:
|
||||
# Generate transpose list
|
||||
transpose_list = []
|
||||
for d in supported_layout:
|
||||
transpose_list.append(layout.index(d))
|
||||
return input_tensor.transpose(transpose_list)
|
||||
return input_tensor
|
||||
|
||||
|
||||
def get_slice(
|
||||
spatial_dimensions: int,
|
||||
pad_np: np.array,
|
||||
dim_coord: tuple[int],
|
||||
kernel: tuple[int],
|
||||
strides: tuple[int],
|
||||
dilation: tuple[int],
|
||||
) -> tuple[slice]:
|
||||
"""
|
||||
Programmatically create a slice object of the right dimensions for pad_np.
|
||||
|
||||
We assume pad_np's first two dimensions are not spatial and are not touched by the pad.
|
||||
|
||||
pad_np[slice] should give the elements of the data that a pool operation will use for the
|
||||
step given in dim_coord.
|
||||
"""
|
||||
slices = [slice(None)] * spatial_dimensions
|
||||
|
||||
for nd in range(spatial_dimensions):
|
||||
slices[nd] = slice(
|
||||
dim_coord[nd] * strides[nd],
|
||||
dim_coord[nd] * strides[nd] + (kernel[nd] - 1) * dilation[nd] + 1,
|
||||
dilation[nd],
|
||||
)
|
||||
|
||||
# Add back batch and channel dimensions
|
||||
slices = [slice(None), slice(None)] + slices
|
||||
|
||||
return tuple(slices)
|
||||
|
||||
|
||||
def pad_tensor(
|
||||
np_arr: np.array,
|
||||
pad_value: float,
|
||||
padding_before: list[int],
|
||||
padding_after: list[int],
|
||||
dtype: str,
|
||||
) -> np.array:
|
||||
"""Pad the spatial dimensions of the given array."""
|
||||
orig_shape = list(np_arr.shape)
|
||||
padded_shape = list(np_arr.shape)
|
||||
n = len(orig_shape)
|
||||
for dim in range(2, n):
|
||||
i = dim - 2
|
||||
padded_shape[dim] += padding_after[i] + padding_before[i]
|
||||
|
||||
pad_np = (np.zeros(shape=padded_shape) + pad_value).astype(dtype)
|
||||
ranges_it = [range(padded_shape[0]), range(padded_shape[1])]
|
||||
for dim in range(2, n):
|
||||
i = dim - 2
|
||||
ranges_it.append(range(padding_before[i], padding_before[i] + orig_shape[dim]))
|
||||
pad_np[np.ix_(*ranges_it)] = np_arr
|
||||
return pad_np
|
||||
|
||||
|
||||
def poolnd_python(
|
||||
np_data: np.array,
|
||||
kernel: tuple[int],
|
||||
strides: tuple[int],
|
||||
dilation: tuple[int],
|
||||
padding_before: tuple[int],
|
||||
padding_after: tuple[int],
|
||||
pool_type: str,
|
||||
count_include_pad: bool = True,
|
||||
ceil_mode: bool = False,
|
||||
dtype: str = "float32",
|
||||
layout: str | None = None,
|
||||
) -> np.array:
|
||||
"""Ground truth pooling operator impelmented in numpy."""
|
||||
|
||||
np_data = _convert_from_layout(np_data, layout)
|
||||
|
||||
out_shape = [np_data.shape[0], np_data.shape[1]]
|
||||
for dim in range(2, len(np_data.shape)):
|
||||
i = dim - 2
|
||||
val = (
|
||||
float(
|
||||
np_data.shape[dim]
|
||||
- (kernel[i] - 1) * dilation[i]
|
||||
- 1
|
||||
+ padding_before[i]
|
||||
+ padding_after[i]
|
||||
)
|
||||
/ strides[i]
|
||||
)
|
||||
|
||||
if ceil_mode:
|
||||
out_shape.append(int(math.ceil(val) + 1))
|
||||
else:
|
||||
out_shape.append(int(math.floor(val) + 1))
|
||||
out_shape = tuple(out_shape)
|
||||
|
||||
# Create a padded array, and a boolean mask showing which values are padded values
|
||||
pad_value = 0
|
||||
if pool_type == "max" and not count_include_pad:
|
||||
pad_value = tvm.te.min_value(dtype).value
|
||||
pad_data = pad_tensor(np_data, pad_value, padding_before, padding_after, dtype)
|
||||
pad_map = pad_tensor(np.ones_like(np_data), 0, padding_before, padding_after, "bool")
|
||||
|
||||
# Create iterator which gives all indices for output array
|
||||
dim_iterators = []
|
||||
for spatial_dimension in range(2, len(np_data.shape)):
|
||||
dim_iterators.append(range(out_shape[spatial_dimension]))
|
||||
coord_iterator = itertools.product(*dim_iterators)
|
||||
|
||||
ret_np = np.zeros(shape=out_shape).astype(dtype)
|
||||
for coordinate in coord_iterator:
|
||||
# Get index into the values that any pool operation will use for given coordinate
|
||||
np_index = get_slice(
|
||||
spatial_dimensions=len(out_shape) - 2,
|
||||
pad_np=pad_data,
|
||||
dim_coord=coordinate,
|
||||
kernel=kernel,
|
||||
strides=strides,
|
||||
dilation=dilation,
|
||||
)
|
||||
|
||||
output_slice = (slice(None), slice(None)) + tuple(coordinate)
|
||||
reduction_axis = tuple(range(2, len(np_data.shape)))
|
||||
if pool_type == "avg":
|
||||
count_non_padded = (
|
||||
pad_data[np_index].size if count_include_pad else np.sum(pad_map[np_index])
|
||||
)
|
||||
# We summed over the non spatial dimensions too so divide by them
|
||||
count_non_padded /= out_shape[0] * out_shape[1]
|
||||
if count_non_padded == 0:
|
||||
ret_np[output_slice] = 0
|
||||
else:
|
||||
ret_np[output_slice] = (
|
||||
np.sum(pad_data[np_index], axis=reduction_axis) / count_non_padded
|
||||
)
|
||||
elif pool_type == "max":
|
||||
count_non_padded = np.sum(pad_map[np_index])
|
||||
# All padded values, default to 0
|
||||
ret_np[output_slice] = np.max(pad_data[np_index], axis=reduction_axis)
|
||||
else:
|
||||
raise ValueError(f"Pool type {pool_type} is not supported")
|
||||
|
||||
return _convert_to_layout(ret_np, layout)
|
||||
@@ -0,0 +1,62 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Reorg in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def reorg_python(a_np, stride):
|
||||
"""Reorg operator
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
stride : int
|
||||
Stride size
|
||||
|
||||
Returns
|
||||
-------
|
||||
b_np : np.ndarray
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
|
||||
batch, in_channel, in_height, in_width = a_np.shape
|
||||
a_np = np.reshape(a_np, batch * in_channel * in_height * in_width)
|
||||
out_c = int(in_channel / (stride * stride))
|
||||
out_channel = in_channel * stride * stride
|
||||
out_height = int(in_height / stride)
|
||||
out_width = int(in_width / stride)
|
||||
b_np = np.zeros(batch * out_channel * out_height * out_width)
|
||||
cnt = 0
|
||||
for b in range(batch):
|
||||
for k in range(in_channel):
|
||||
for j in range(in_height):
|
||||
for i in range(in_width):
|
||||
c2 = k % out_c
|
||||
offset = int(k / out_c)
|
||||
w2 = int(i * stride + offset % stride)
|
||||
h2 = int(j * stride + offset / stride)
|
||||
out_index = int(
|
||||
w2 + in_width * stride * (h2 + in_height * stride * (c2 + out_c * b))
|
||||
)
|
||||
b_np[cnt] = a_np[int(out_index)]
|
||||
cnt = cnt + 1
|
||||
b_np = np.reshape(b_np, (batch, out_channel, out_height, out_width))
|
||||
return b_np
|
||||
@@ -0,0 +1,307 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
# ruff: noqa: E741, F841, RUF005
|
||||
"""Upsampling in python"""
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.topi.utils import nchw_pack_layout
|
||||
|
||||
|
||||
def get_inx(x, image_width, target_width, coordinate_transformation_mode):
|
||||
"""Infer input x from output x with various coordinate transformation methods"""
|
||||
scale = image_width / target_width
|
||||
if coordinate_transformation_mode == "half_pixel":
|
||||
in_x = (x + 0.5) * scale - 0.5
|
||||
elif coordinate_transformation_mode == "align_corners":
|
||||
in_x = (image_width - 1) / (target_width - 1) * x if target_width > 1 else 0
|
||||
elif coordinate_transformation_mode == "asymmetric":
|
||||
in_x = scale * x
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported coordinate_transformation_mode: {coordinate_transformation_mode}"
|
||||
)
|
||||
return in_x
|
||||
|
||||
|
||||
def get_index(x, image_width, target_width, coordinate_transformation_mode, rounding_method=""):
|
||||
"""get and round the nearest index for nearest_neighbor"""
|
||||
in_x = get_inx(x, image_width, target_width, coordinate_transformation_mode)
|
||||
|
||||
effective_rounding_method = rounding_method
|
||||
if not effective_rounding_method:
|
||||
if coordinate_transformation_mode == "align_corners":
|
||||
effective_rounding_method = "round"
|
||||
else:
|
||||
effective_rounding_method = "floor"
|
||||
|
||||
if effective_rounding_method == "floor":
|
||||
out = math.floor(in_x)
|
||||
elif effective_rounding_method == "round":
|
||||
out = round(in_x)
|
||||
elif effective_rounding_method == "round_prefer_floor":
|
||||
out = math.ceil(in_x - 0.5)
|
||||
elif effective_rounding_method == "round_prefer_ceil":
|
||||
out = math.floor(in_x + 0.5)
|
||||
elif effective_rounding_method == "ceil":
|
||||
out = math.ceil(in_x)
|
||||
else:
|
||||
raise ValueError(f"Unknown rounding method: {rounding_method!r}")
|
||||
|
||||
out = max(min(out, image_width - 1), 0)
|
||||
return int(out)
|
||||
|
||||
|
||||
def resize3d_nearest(arr, scale, coordinate_transformation_mode, rounding_method=""):
|
||||
"""Populate the array by scale factor"""
|
||||
d, h, w = arr.shape
|
||||
out_d, out_h, out_w = [round(i * s) for i, s in zip(arr.shape, scale)]
|
||||
out = np.empty((out_d, out_h, out_w))
|
||||
for z in range(out_d):
|
||||
for y in range(out_h):
|
||||
for x in range(out_w):
|
||||
in_z = get_index(z, d, out_d, coordinate_transformation_mode, rounding_method)
|
||||
in_y = get_index(y, h, out_h, coordinate_transformation_mode, rounding_method)
|
||||
in_x = get_index(x, w, out_w, coordinate_transformation_mode, rounding_method)
|
||||
out[z, y, x] = arr[in_z, in_y, in_x]
|
||||
return out
|
||||
|
||||
|
||||
def resize3d_linear(data_in, scale, coordinate_transformation_mode):
|
||||
"""Trilinear 3d scaling using python"""
|
||||
dtype = data_in.dtype
|
||||
d, h, w = data_in.shape
|
||||
new_d, new_h, new_w = [round(i * s) for i, s in zip(data_in.shape, scale)]
|
||||
data_out = np.ones((new_d, new_h, new_w))
|
||||
|
||||
indexes = np.mgrid[0:2, 0:2, 0:2]
|
||||
|
||||
def _get_patch(zint, yint, xint):
|
||||
# Get the surrounding values
|
||||
indices = indexes.copy()
|
||||
indices[0] = np.maximum(np.minimum(indexes[0] + zint, d - 1), 0)
|
||||
indices[1] = np.maximum(np.minimum(indexes[1] + yint, h - 1), 0)
|
||||
indices[2] = np.maximum(np.minimum(indexes[2] + xint, w - 1), 0)
|
||||
p = data_in[indices[0], indices[1], indices[2]]
|
||||
return p
|
||||
|
||||
for m in range(new_d):
|
||||
for j in range(new_h):
|
||||
for k in range(new_w):
|
||||
in_z = get_inx(m, d, new_d, coordinate_transformation_mode)
|
||||
in_y = get_inx(j, h, new_h, coordinate_transformation_mode)
|
||||
in_x = get_inx(k, w, new_w, coordinate_transformation_mode)
|
||||
zint = math.floor(in_z)
|
||||
zfract = in_z - math.floor(in_z)
|
||||
|
||||
yint = math.floor(in_y)
|
||||
yfract = in_y - math.floor(in_y)
|
||||
|
||||
xint = math.floor(in_x)
|
||||
xfract = in_x - math.floor(in_x)
|
||||
|
||||
wz = np.array([1.0 - zfract, zfract], dtype=dtype)
|
||||
wy = np.array([1.0 - yfract, yfract], dtype=dtype)
|
||||
wx = np.array([1.0 - xfract, xfract], dtype=dtype)
|
||||
|
||||
p = _get_patch(zint, yint, xint)
|
||||
l = np.sum(p * wx, axis=-1)
|
||||
col = np.sum(l * wy, axis=-1)
|
||||
data_out[m, j, k] = np.sum(col * wz)
|
||||
|
||||
return data_out
|
||||
|
||||
|
||||
def resize3d_cubic(data_in, scale, coordinate_transformation_mode):
|
||||
"""Tricubic 3d scaling using python"""
|
||||
dtype = data_in.dtype
|
||||
d, h, w = data_in.shape
|
||||
new_d, new_h, new_w = [round(i * s) for i, s in zip(data_in.shape, scale)]
|
||||
data_out = np.ones((new_d, new_h, new_w))
|
||||
|
||||
def _cubic_spline_weights(t, alpha=-0.5):
|
||||
"""create cubic spline weights in 1D"""
|
||||
t2 = t * t
|
||||
t3 = t * t * t
|
||||
w1 = alpha * (t3 - 2 * t2 + t)
|
||||
w2 = (alpha + 2) * t3 - (3 + alpha) * t2 + 1
|
||||
w3 = -(alpha + 2) * t3 + (3 + 2 * alpha) * t2 - alpha * t
|
||||
w4 = -alpha * t3 + alpha * t2
|
||||
return np.array([w1, w2, w3, w4])
|
||||
|
||||
indexes = np.mgrid[-1:3, -1:3, -1:3]
|
||||
|
||||
def _get_patch(zint, yint, xint):
|
||||
# Get the surrounding values
|
||||
indices = indexes.copy()
|
||||
indices[0] = np.maximum(np.minimum(indexes[0] + zint, d - 1), 0)
|
||||
indices[1] = np.maximum(np.minimum(indexes[1] + yint, h - 1), 0)
|
||||
indices[2] = np.maximum(np.minimum(indexes[2] + xint, w - 1), 0)
|
||||
p = data_in[indices[0], indices[1], indices[2]]
|
||||
return p
|
||||
|
||||
for m in range(new_d):
|
||||
for j in range(new_h):
|
||||
for k in range(new_w):
|
||||
in_z = get_inx(m, d, new_d, coordinate_transformation_mode)
|
||||
in_y = get_inx(j, h, new_h, coordinate_transformation_mode)
|
||||
in_x = get_inx(k, w, new_w, coordinate_transformation_mode)
|
||||
zint = math.floor(in_z)
|
||||
zfract = in_z - math.floor(in_z)
|
||||
|
||||
yint = math.floor(in_y)
|
||||
yfract = in_y - math.floor(in_y)
|
||||
|
||||
xint = math.floor(in_x)
|
||||
xfract = in_x - math.floor(in_x)
|
||||
|
||||
wz = _cubic_spline_weights(zfract)
|
||||
wy = _cubic_spline_weights(yfract)
|
||||
wx = _cubic_spline_weights(xfract)
|
||||
|
||||
p = _get_patch(zint, yint, xint)
|
||||
|
||||
l = np.sum(p * wx, axis=-1)
|
||||
col = np.sum(l * wy, axis=-1)
|
||||
data_out[m, j, k] = np.sum(col * wz)
|
||||
|
||||
return data_out
|
||||
|
||||
|
||||
def resize3d_ncdhw(
|
||||
data,
|
||||
scale,
|
||||
method="nearest_neighbor",
|
||||
coordinate_transformation_mode="align_corners",
|
||||
rounding_method="",
|
||||
):
|
||||
"""reference kernel for 3D image resizing"""
|
||||
ishape = data.shape
|
||||
|
||||
oshape = (
|
||||
ishape[0],
|
||||
ishape[1],
|
||||
round(ishape[2] * scale[0]),
|
||||
round(ishape[3] * scale[1]),
|
||||
round(ishape[4] * scale[2]),
|
||||
)
|
||||
|
||||
output_np = np.zeros(oshape, dtype=data.dtype)
|
||||
|
||||
for b in range(oshape[0]):
|
||||
for c in range(oshape[1]):
|
||||
if method == "nearest_neighbor":
|
||||
output_np[b, c, :, :, :] = resize3d_nearest(
|
||||
data[b, c, :, :, :], scale, coordinate_transformation_mode, rounding_method
|
||||
)
|
||||
elif method == "linear":
|
||||
output_np[b, c, :, :, :] = resize3d_linear(
|
||||
data[b, c, :, :, :], scale, coordinate_transformation_mode
|
||||
)
|
||||
elif method == "cubic":
|
||||
output_np[b, c, :, :, :] = resize3d_cubic(
|
||||
data[b, c, :, :, :], scale, coordinate_transformation_mode
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unknown resize method", method)
|
||||
|
||||
return output_np
|
||||
|
||||
|
||||
def resize1d_python(
|
||||
data,
|
||||
scale,
|
||||
layout="NCW",
|
||||
method="nearest_neighbor",
|
||||
coordinate_transformation_mode="align_corners",
|
||||
rounding_method="",
|
||||
):
|
||||
"""Python version of 3D scaling using nearest neighbour"""
|
||||
|
||||
if layout == "NWC":
|
||||
data = data.transpose([0, 2, 1])
|
||||
|
||||
data = np.expand_dims(data, axis=[2, 3])
|
||||
output_np = resize3d_ncdhw(
|
||||
data, (1, 1) + scale, method, coordinate_transformation_mode, rounding_method
|
||||
)
|
||||
output_np = np.squeeze(output_np, axis=2)
|
||||
output_np = np.squeeze(output_np, axis=2)
|
||||
|
||||
if layout == "NWC":
|
||||
output_np = output_np.transpose([0, 2, 1])
|
||||
|
||||
return output_np
|
||||
|
||||
|
||||
def resize2d_python(
|
||||
data,
|
||||
scale,
|
||||
layout="NCHW",
|
||||
method="nearest_neighbor",
|
||||
coordinate_transformation_mode="align_corners",
|
||||
rounding_method="",
|
||||
):
|
||||
"""Python version of scaling using nearest neighbour"""
|
||||
|
||||
if layout == "NHWC":
|
||||
data = data.transpose([0, 3, 1, 2])
|
||||
elif nchw_pack_layout(layout):
|
||||
ishape = data.shape
|
||||
transposed = data.transpose([0, 4, 1, 5, 2, 3])
|
||||
tshape = transposed.shape
|
||||
data = transposed.reshape(
|
||||
tshape[0] * tshape[1], tshape[2] * tshape[3], tshape[4], tshape[5]
|
||||
)
|
||||
|
||||
data = np.expand_dims(data, axis=2)
|
||||
output_np = resize3d_ncdhw(
|
||||
data, (1,) + scale, method, coordinate_transformation_mode, rounding_method
|
||||
)
|
||||
output_np = np.squeeze(output_np, axis=2)
|
||||
|
||||
if layout == "NHWC":
|
||||
output_np = output_np.transpose([0, 2, 3, 1])
|
||||
elif nchw_pack_layout(layout):
|
||||
output_np = output_np.reshape(tshape[0:4] + output_np.shape[2:])
|
||||
output_np = output_np.transpose([0, 2, 4, 5, 1, 3])
|
||||
|
||||
return output_np
|
||||
|
||||
|
||||
def resize3d_python(
|
||||
data,
|
||||
scale,
|
||||
layout="NCDHW",
|
||||
method="nearest_neighbor",
|
||||
coordinate_transformation_mode="align_corners",
|
||||
rounding_method="",
|
||||
):
|
||||
"""Python version of 3D scaling using nearest neighbour"""
|
||||
|
||||
if layout == "NDHWC":
|
||||
data = data.transpose([0, 4, 1, 2, 3])
|
||||
|
||||
output_np = resize3d_ncdhw(data, scale, method, coordinate_transformation_mode, rounding_method)
|
||||
|
||||
if layout == "NDHWC":
|
||||
output_np = output_np.transpose([0, 2, 3, 4, 1])
|
||||
|
||||
return output_np
|
||||
@@ -0,0 +1,53 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Root mean square normalization in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def rms_norm_python(data, weight, axis, epsilon=1e-5):
|
||||
"""Root mean square normalization operator in Python.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
weight: numpy.ndarray
|
||||
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
bias: numpy.ndarray
|
||||
Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
axis : int or tuple of ints
|
||||
Axis over the normalization applied
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : np.ndarray
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
dtype = data.dtype
|
||||
data = data.astype("float32")
|
||||
weight = weight.astype("float32")
|
||||
square_mean = np.mean(np.square(data), axis, keepdims=True)
|
||||
result = data * weight / np.sqrt(square_mean + epsilon)
|
||||
return result.astype(dtype)
|
||||
@@ -0,0 +1,184 @@
|
||||
# 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, too-many-nested-blocks
|
||||
"Roi align in python"
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _bilinear(a_np, n, c, y, x, height, width, layout):
|
||||
if y < -1 or y > height or x < -1 or x > width:
|
||||
return 0
|
||||
|
||||
y = min(max(y, 0), height - 1)
|
||||
x = min(max(x, 0), width - 1)
|
||||
|
||||
y_low = math.floor(y)
|
||||
x_low = math.floor(x)
|
||||
y_high = y_low + 1
|
||||
x_high = x_low + 1
|
||||
|
||||
wy_h = y - y_low
|
||||
wx_h = x - x_low
|
||||
wy_l = 1 - wy_h
|
||||
wx_l = 1 - wx_h
|
||||
|
||||
val = 0
|
||||
for wx, xp in zip((wx_l, wx_h), (x_low, x_high)):
|
||||
for wy, yp in zip((wy_l, wy_h), (y_low, y_high)):
|
||||
if 0 <= yp < height and 0 <= xp < width:
|
||||
if layout == "NCHW":
|
||||
val += wx * wy * a_np[n, c, yp, xp]
|
||||
else:
|
||||
val += wx * wy * a_np[n, yp, xp, c]
|
||||
return val
|
||||
|
||||
|
||||
def roi_align_common(
|
||||
a_np,
|
||||
b_np,
|
||||
rois_np,
|
||||
channel,
|
||||
pooled_size_h,
|
||||
pooled_size_w,
|
||||
spatial_scale,
|
||||
sample_ratio,
|
||||
aligned,
|
||||
avg_mode,
|
||||
max_mode,
|
||||
height,
|
||||
width,
|
||||
layout,
|
||||
):
|
||||
"""Common code used by roi align NCHW and NHWC"""
|
||||
num_roi = rois_np.shape[0]
|
||||
|
||||
for i in range(num_roi):
|
||||
roi = rois_np[i]
|
||||
batch_index = int(roi[0])
|
||||
roi_start_w, roi_start_h, roi_end_w, roi_end_h = roi[1:] * spatial_scale
|
||||
roi_h = roi_end_h - roi_start_h if aligned else max(roi_end_h - roi_start_h, 1.0)
|
||||
roi_w = roi_end_w - roi_start_w if aligned else max(roi_end_w - roi_start_w, 1.0)
|
||||
|
||||
bin_h = roi_h / pooled_size_h
|
||||
bin_w = roi_w / pooled_size_w
|
||||
|
||||
if sample_ratio > 0:
|
||||
roi_bin_grid_h = roi_bin_grid_w = int(sample_ratio)
|
||||
else:
|
||||
roi_bin_grid_h = math.ceil(roi_h / pooled_size_h)
|
||||
roi_bin_grid_w = math.ceil(roi_w / pooled_size_w)
|
||||
|
||||
count = roi_bin_grid_h * roi_bin_grid_w
|
||||
|
||||
for c in range(channel):
|
||||
for ph in range(pooled_size_h):
|
||||
for pw in range(pooled_size_w):
|
||||
if avg_mode:
|
||||
total = 0.0
|
||||
if max_mode:
|
||||
total = float("-inf")
|
||||
for iy in range(roi_bin_grid_h):
|
||||
for ix in range(roi_bin_grid_w):
|
||||
y = roi_start_h + ph * bin_h + (iy + 0.5) * bin_h / roi_bin_grid_h
|
||||
x = roi_start_w + pw * bin_w + (ix + 0.5) * bin_w / roi_bin_grid_w
|
||||
if avg_mode:
|
||||
total += (
|
||||
_bilinear(a_np, batch_index, c, y, x, height, width, layout)
|
||||
/ count
|
||||
)
|
||||
if max_mode:
|
||||
total = max(
|
||||
total,
|
||||
_bilinear(a_np, batch_index, c, y, x, height, width, layout),
|
||||
)
|
||||
|
||||
if layout == "NCHW":
|
||||
b_np[i, c, ph, pw] = total
|
||||
else:
|
||||
b_np[i, ph, pw, c] = total
|
||||
return b_np
|
||||
|
||||
|
||||
def roi_align_nchw_python(
|
||||
a_np, rois_np, pooled_size, spatial_scale, sample_ratio, mode=b"avg", aligned=False
|
||||
):
|
||||
"""Roi align NCHW in python"""
|
||||
avg_mode = mode in (b"avg", "avg", 0)
|
||||
max_mode = mode in (b"max", "max", 1)
|
||||
assert avg_mode or max_mode, "Mode must be average or max. Please pass a valid mode."
|
||||
_, channel, height, width = a_np.shape
|
||||
if isinstance(pooled_size, int):
|
||||
pooled_size_h = pooled_size_w = pooled_size
|
||||
else:
|
||||
pooled_size_h, pooled_size_w = pooled_size
|
||||
|
||||
b_np = np.zeros((rois_np.shape[0], channel, pooled_size_h, pooled_size_w), dtype=a_np.dtype)
|
||||
|
||||
return roi_align_common(
|
||||
a_np,
|
||||
b_np,
|
||||
rois_np,
|
||||
channel,
|
||||
pooled_size_h,
|
||||
pooled_size_w,
|
||||
spatial_scale,
|
||||
sample_ratio,
|
||||
aligned,
|
||||
avg_mode,
|
||||
max_mode,
|
||||
height,
|
||||
width,
|
||||
"NCHW",
|
||||
)
|
||||
|
||||
|
||||
def roi_align_nhwc_python(
|
||||
a_np, rois_np, pooled_size, spatial_scale, sample_ratio, mode=b"avg", aligned=False
|
||||
):
|
||||
"""Roi align NHWC in python"""
|
||||
avg_mode = mode in (b"avg", "avg", 0)
|
||||
max_mode = mode in (b"max", "max", 1)
|
||||
assert avg_mode or max_mode, "Mode must be average or max. Please pass a valid mode."
|
||||
_, height, width, channel = a_np.shape
|
||||
num_roi = rois_np.shape[0]
|
||||
|
||||
if isinstance(pooled_size, int):
|
||||
pooled_size_h = pooled_size_w = pooled_size
|
||||
else:
|
||||
pooled_size_h, pooled_size_w = pooled_size
|
||||
|
||||
b_np = np.zeros((num_roi, pooled_size_h, pooled_size_w, channel), dtype=a_np.dtype)
|
||||
|
||||
return roi_align_common(
|
||||
a_np,
|
||||
b_np,
|
||||
rois_np,
|
||||
channel,
|
||||
pooled_size_h,
|
||||
pooled_size_w,
|
||||
spatial_scale,
|
||||
sample_ratio,
|
||||
aligned,
|
||||
avg_mode,
|
||||
max_mode,
|
||||
height,
|
||||
width,
|
||||
"NHWC",
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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, too-many-nested-blocks
|
||||
"Roi pool in python"
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def roi_pool_nchw_python(a_np, rois_np, pooled_size, spatial_scale):
|
||||
"""Roi pool in python"""
|
||||
_, channel, height, width = a_np.shape
|
||||
num_roi = rois_np.shape[0]
|
||||
b_np = np.zeros((num_roi, channel, pooled_size, pooled_size), dtype=a_np.dtype)
|
||||
|
||||
if isinstance(pooled_size, int):
|
||||
pooled_size_h = pooled_size_w = pooled_size
|
||||
else:
|
||||
pooled_size_h, pooled_size_w = pooled_size
|
||||
|
||||
for i in range(num_roi):
|
||||
roi = rois_np[i]
|
||||
batch_index = int(roi[0])
|
||||
# Use ties-away-from-zero rounding to match ONNX runtime (std::round semantics).
|
||||
# Python's built-in round() uses ties-to-even, so use floor(x + 0.5) explicitly.
|
||||
roi_start_w = math.floor(roi[1] * spatial_scale + 0.5)
|
||||
roi_start_h = math.floor(roi[2] * spatial_scale + 0.5)
|
||||
roi_end_w = math.floor(roi[3] * spatial_scale + 0.5)
|
||||
roi_end_h = math.floor(roi[4] * spatial_scale + 0.5)
|
||||
roi_h = max(roi_end_h - roi_start_h + 1, 1)
|
||||
roi_w = max(roi_end_w - roi_start_w + 1, 1)
|
||||
|
||||
bin_h = float(roi_h) / pooled_size_h
|
||||
bin_w = float(roi_w) / pooled_size_w
|
||||
|
||||
for ph in range(pooled_size_h):
|
||||
for pw in range(pooled_size_w):
|
||||
hstart = math.floor(ph * bin_h)
|
||||
wstart = math.floor(pw * bin_w)
|
||||
hend = math.ceil((ph + 1) * bin_h)
|
||||
wend = math.ceil((pw + 1) * bin_w)
|
||||
hstart = min(max(hstart + roi_start_h, 0), height)
|
||||
hend = min(max(hend + roi_start_h, 0), height)
|
||||
wstart = min(max(wstart + roi_start_w, 0), width)
|
||||
wend = min(max(wend + roi_start_w, 0), width)
|
||||
is_empty = (hend <= hstart) or (wend <= wstart)
|
||||
|
||||
for c in range(channel):
|
||||
if is_empty:
|
||||
b_np[i, c, ph, pw] = 0.0
|
||||
else:
|
||||
b_np[i, c, ph, pw] = np.max(a_np[batch_index, c, hstart:hend, wstart:wend])
|
||||
return b_np
|
||||
@@ -0,0 +1,36 @@
|
||||
# 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.
|
||||
"""The reference implementation of searchsorted in Numpy."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def searchsorted_ref(sorted_sequence, values, right, out_dtype):
|
||||
"""Run Numpy searchsorted on 1-D or N-D sorted_sequence."""
|
||||
side = "right" if right else "left"
|
||||
if len(sorted_sequence.shape) == 1 and len(values.shape) > 1:
|
||||
sorted_sequence_2d = np.tile(sorted_sequence, (np.prod(values.shape[:-1]), 1))
|
||||
else:
|
||||
sorted_sequence_2d = np.reshape(sorted_sequence, (-1, sorted_sequence.shape[-1]))
|
||||
|
||||
values_2d = np.reshape(values, (-1, values.shape[-1]))
|
||||
indices = np.zeros(values_2d.shape, dtype=out_dtype)
|
||||
|
||||
for i in range(indices.shape[0]):
|
||||
indices[i] = np.searchsorted(sorted_sequence_2d[i], values_2d[i], side=side)
|
||||
|
||||
return np.reshape(indices, values.shape)
|
||||
@@ -0,0 +1,55 @@
|
||||
# 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
|
||||
"""Sequence mask in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def sequence_mask(data, valid_length, mask_value, axis):
|
||||
"""batch_matmul operator implemented in numpy.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
N-D with shape [batch_size, MAX_LENGTH, ...] or [MAX_LENGTH, batch_size, ...]
|
||||
|
||||
valid_length : numpy.ndarray
|
||||
1-D with shape [batch_size,]
|
||||
|
||||
mask_value : float
|
||||
Masking value
|
||||
|
||||
axis : int
|
||||
The axis of the length dimension
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : numpy.ndarray
|
||||
N-D with shape same as data
|
||||
"""
|
||||
in_shape = data.shape
|
||||
max_length = data.shape[axis]
|
||||
val_len_expand_shape = [1 for _ in range(len(in_shape))]
|
||||
val_len_expand_shape[1 - axis] = in_shape[1 - axis]
|
||||
seq_len_expand_shape = [1 for _ in range(len(in_shape))]
|
||||
seq_len_expand_shape[axis] = in_shape[axis]
|
||||
mask = np.broadcast_to(
|
||||
np.arange(max_length).reshape(seq_len_expand_shape), in_shape
|
||||
) >= valid_length.reshape(val_len_expand_shape)
|
||||
out = data * (1 - mask) + mask_value * mask
|
||||
return out
|
||||
@@ -0,0 +1,51 @@
|
||||
# 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.
|
||||
"""Slice axis in python"""
|
||||
|
||||
|
||||
def slice_axis_python(data, axis, begin, end=None):
|
||||
"""Slice input array along specific axis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
The source array to be sliced.
|
||||
|
||||
axis : int
|
||||
Axis to be sliced.
|
||||
|
||||
begin: int
|
||||
The index to begin with in the slicing.
|
||||
|
||||
end: int, optional
|
||||
The index indicating end of the slice.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : numpy.ndarray
|
||||
The computed result.
|
||||
"""
|
||||
dshape = data.shape
|
||||
if axis < 0:
|
||||
axis += len(dshape)
|
||||
if begin < 0:
|
||||
begin += dshape[axis]
|
||||
if end <= 0:
|
||||
end += dshape[axis]
|
||||
slc = [slice(None)] * len(dshape)
|
||||
slc[axis] = slice(begin, end)
|
||||
return data[tuple(slc)]
|
||||
@@ -0,0 +1,58 @@
|
||||
# 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, trailing-whitespace
|
||||
"""Softmax and log_softmax operation in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def softmax_python(a_np, axis=1):
|
||||
"""Softmax operator.
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
N-D input data
|
||||
|
||||
Returns
|
||||
-------
|
||||
output_np : numpy.ndarray
|
||||
N-D output with same shape
|
||||
"""
|
||||
max_elem = np.amax(a_np, axis=axis, keepdims=True)
|
||||
e = np.exp(a_np - max_elem)
|
||||
expsum = np.sum(e, axis=axis, keepdims=True)
|
||||
out_np = e / expsum
|
||||
return out_np
|
||||
|
||||
|
||||
def log_softmax_python(a_np, axis=1):
|
||||
"""Log_softmax operator.
|
||||
Parameters
|
||||
----------
|
||||
a_np : numpy.ndarray
|
||||
N-D input data
|
||||
|
||||
Returns
|
||||
-------
|
||||
output_np : numpy.ndarray
|
||||
N-D output with same shape
|
||||
"""
|
||||
max_elem = np.amax(a_np, axis=axis, keepdims=True)
|
||||
e = np.exp(a_np - max_elem)
|
||||
expsum = np.sum(e, axis=axis, keepdims=True)
|
||||
out_np = a_np - max_elem - np.log(expsum)
|
||||
return out_np
|
||||
@@ -0,0 +1,94 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Space to batch ND in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def space_to_batch_nd_python(data, block_shape, pad_before, pad_after, pad_value=0):
|
||||
"""Space to Batch operator in python for NHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : np.ndarray
|
||||
N-D with shape [batch, spatial_shape, remaining_shapes],
|
||||
where spatial_shape has M dimensions.
|
||||
|
||||
block_shape : list of ints
|
||||
1-D array of size [M] where M is number of spatial dims, specifies block
|
||||
size for each spatial dimension.
|
||||
|
||||
pad_before : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
zero-padding size before each spatial dimension.
|
||||
|
||||
pad_after : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
zero-padding size after each spatial dimension.
|
||||
|
||||
pad_value : float, optional
|
||||
the value used for padding. Defaults to 0.
|
||||
|
||||
Returns
|
||||
-------
|
||||
s2b_out : np.ndarray
|
||||
N-D with shape [batch * prod(block_shape),
|
||||
padded_data[1] / block_shape[0], ..., padded_data[M] / block_shape[M-1],
|
||||
remaining_shape]
|
||||
"""
|
||||
M = len(block_shape)
|
||||
in_batch = data.shape[0]
|
||||
block_shape_prod = np.prod(block_shape)
|
||||
|
||||
# Apply padding to input data
|
||||
input_shape = data.shape
|
||||
# Add the paddings for batch and remaining dims
|
||||
paddings = map(list, zip(pad_before, pad_after))
|
||||
paddings = [[0, 0]] + list(paddings) + [[0, 0]] * (data.ndim - 1 - M)
|
||||
padded_data = np.pad(data, paddings, mode="constant", constant_values=pad_value)
|
||||
padded_shape = padded_data.shape
|
||||
|
||||
# Get the reshape shape and transpose axes
|
||||
r_shape = []
|
||||
trans_axis = []
|
||||
r_shape.append(in_batch)
|
||||
for i in range(1, M + 1):
|
||||
r_shape.append(int(padded_shape[i] // block_shape[i - 1]))
|
||||
r_shape.append(block_shape[i - 1])
|
||||
trans_axis.append(len(r_shape) - 1)
|
||||
|
||||
axis_len = len(trans_axis)
|
||||
trans_axis.append(0)
|
||||
for i in range(axis_len):
|
||||
trans_axis.append(trans_axis[i] - 1)
|
||||
|
||||
out_shape = []
|
||||
out_shape.append(int(in_batch * block_shape_prod))
|
||||
for i in range(1, M + 1):
|
||||
out_shape.append(int(padded_shape[i] // block_shape[i - 1]))
|
||||
|
||||
for i in range(M + 1, len(input_shape)):
|
||||
r_shape.append(input_shape[i])
|
||||
trans_axis.append(len(r_shape) - 1)
|
||||
out_shape.append(input_shape[i])
|
||||
|
||||
s2b_out = np.reshape(padded_data, newshape=r_shape)
|
||||
s2b_out = np.transpose(s2b_out, axes=trans_axis)
|
||||
s2b_out = np.reshape(s2b_out, newshape=out_shape)
|
||||
|
||||
return s2b_out
|
||||
@@ -0,0 +1,49 @@
|
||||
# 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, line-too-long, unused-variable, too-many-locals
|
||||
"""Space to depth in python"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def space_to_depth_python(data, block_size):
|
||||
"""Space to Depth operator in python for NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : np.ndarray
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
block_size : int
|
||||
Size of spatial blocks to decompose into channels.
|
||||
|
||||
Returns
|
||||
-------
|
||||
d2s_out : np.ndarray
|
||||
4-D with shape [batch, in_channel * (block_size * block_size),
|
||||
out_height / block_size, out_width / block_size]
|
||||
"""
|
||||
in_n, in_c, in_h, in_w = data.shape
|
||||
new_h = int(in_h / block_size)
|
||||
new_w = int(in_h / block_size)
|
||||
new_c = int(in_c * (block_size * block_size))
|
||||
|
||||
expanded = np.reshape(data, newshape=[in_n, in_c, new_h, block_size, new_w, block_size])
|
||||
transposed = np.transpose(expanded, axes=[0, 3, 5, 1, 2, 4])
|
||||
newshape = [in_n, new_c, new_h, new_w]
|
||||
d2s_out = np.reshape(transposed, newshape=newshape)
|
||||
return d2s_out
|
||||
@@ -0,0 +1,129 @@
|
||||
# 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.
|
||||
"""strided_slice/set in python"""
|
||||
|
||||
|
||||
def strided_slice_python(data, begin, end, strides, slice_mode="end", axes=None):
|
||||
"""Python version of strided slice operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
Input data
|
||||
|
||||
begin : list
|
||||
Beginning of the slices.
|
||||
|
||||
end : list
|
||||
End of the slices.
|
||||
|
||||
strides : list
|
||||
The stride of each slice.
|
||||
|
||||
slice_mode : str, optional
|
||||
The slice mode [end, size].
|
||||
|
||||
- ``"end"``: The default slice mode, ending indices for the slice.
|
||||
- ``"size"``: The input strides will be ignored, input end in this mode indicates
|
||||
the size of a slice starting at the location specified by begin. If end[i] is -1,
|
||||
all remaining elements in that dimension are included in the slice.
|
||||
|
||||
axes : list, optional
|
||||
Axes along which slicing is applied
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : numpy.ndarray
|
||||
The sliced result.
|
||||
"""
|
||||
strides = [] if strides is None else strides
|
||||
if axes is not None:
|
||||
rank = len(data.shape)
|
||||
new_begin = [0] * rank
|
||||
new_end = [data.shape[i] for i in range(rank)]
|
||||
new_strides = [1] * rank
|
||||
|
||||
for i, axis in enumerate(axes):
|
||||
new_begin[axis] = begin[i]
|
||||
new_end[axis] = end[i]
|
||||
if len(strides) > i:
|
||||
new_strides[axis] = strides[i]
|
||||
|
||||
begin = new_begin
|
||||
end = new_end
|
||||
strides = new_strides
|
||||
|
||||
slices = []
|
||||
for i in range(len(data.shape)):
|
||||
new_stride = None
|
||||
if slice_mode == "end" and i < len(strides):
|
||||
new_stride = strides[i]
|
||||
|
||||
new_begin = begin[i] if i < len(begin) else None
|
||||
if i >= len(end):
|
||||
new_end = None
|
||||
elif slice_mode == "size":
|
||||
if end[i] < 0:
|
||||
new_end = None
|
||||
else:
|
||||
new_end = new_begin + end[i]
|
||||
else:
|
||||
new_end = end[i]
|
||||
|
||||
slices.append(slice(new_begin, new_end, new_stride))
|
||||
|
||||
return data[tuple(slices)]
|
||||
|
||||
|
||||
def strided_set_python(data, v, begin, end, strides):
|
||||
"""Python version of strided slice operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : numpy.ndarray
|
||||
Input data
|
||||
|
||||
v : numpy.ndarray
|
||||
Value data
|
||||
|
||||
begin : list
|
||||
Beginning of the slices.
|
||||
|
||||
end : list
|
||||
End of the slices.
|
||||
|
||||
strides : list
|
||||
The stride of each slice.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : numpy.ndarray
|
||||
The updated result.
|
||||
"""
|
||||
strides = [] if strides is None else strides
|
||||
slices = []
|
||||
res = data.copy()
|
||||
for i in range(len(data.shape)):
|
||||
slices.append(
|
||||
slice(
|
||||
begin[i] if i < len(begin) else None,
|
||||
end[i] if i < len(end) else None,
|
||||
strides[i] if i < len(strides) else None,
|
||||
)
|
||||
)
|
||||
res[tuple(slices)] = v
|
||||
return res
|
||||
Reference in New Issue
Block a user