wehub-resource-sync
205 lines
6.5 KiB
Python
205 lines
6.5 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name
|
|
"""ROI Align operator"""
|
|
|
|
import tvm
|
|
from tvm import te
|
|
|
|
from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
|
|
|
|
|
|
def _sample_common(
|
|
i,
|
|
c,
|
|
ph,
|
|
pw,
|
|
rois,
|
|
pooled_size_h,
|
|
pooled_size_w,
|
|
spatial_scale,
|
|
sample_ratio,
|
|
aligned,
|
|
dtype,
|
|
avg_mode,
|
|
bilinear_func,
|
|
):
|
|
roi = rois[i]
|
|
batch_index = roi[0].astype("int32")
|
|
roi_start_w = roi[1] * spatial_scale
|
|
roi_start_h = roi[2] * spatial_scale
|
|
roi_end_w = roi[3] * spatial_scale
|
|
roi_end_h = roi[4] * spatial_scale
|
|
|
|
if aligned:
|
|
roi_h = roi_end_h - roi_start_h
|
|
roi_w = roi_end_w - roi_start_w
|
|
else:
|
|
roi_h = te.max(roi_end_h - roi_start_h, tvm.tirx.const(1.0, dtype))
|
|
roi_w = te.max(roi_end_w - roi_start_w, tvm.tirx.const(1.0, dtype))
|
|
|
|
pooled_size_h_const = tvm.tirx.const(pooled_size_h, dtype)
|
|
pooled_size_w_const = tvm.tirx.const(pooled_size_w, dtype)
|
|
bin_h = roi_h / pooled_size_h_const
|
|
bin_w = roi_w / pooled_size_w_const
|
|
|
|
if sample_ratio > 0:
|
|
roi_bin_grid_h = tvm.tirx.const(sample_ratio, "int32")
|
|
roi_bin_grid_w = tvm.tirx.const(sample_ratio, "int32")
|
|
else:
|
|
roi_bin_grid_h = te.ceil(roi_h / pooled_size_h_const).astype("int32")
|
|
roi_bin_grid_w = te.ceil(roi_w / pooled_size_w_const).astype("int32")
|
|
|
|
count = roi_bin_grid_h * roi_bin_grid_w
|
|
rh = te.reduce_axis((0, roi_bin_grid_h), name="rh")
|
|
rw = te.reduce_axis((0, roi_bin_grid_w), name="rw")
|
|
roi_start_h = roi_start_h + tvm.tirx.Cast(dtype, ph) * bin_h
|
|
roi_start_w = roi_start_w + tvm.tirx.Cast(dtype, pw) * bin_w
|
|
|
|
def sample_value(rh_idx, rw_idx):
|
|
return bilinear_func(
|
|
batch_index,
|
|
c,
|
|
roi_start_h
|
|
+ (tvm.tirx.Cast(dtype, rh_idx) + tvm.tirx.const(0.5, dtype))
|
|
* bin_h
|
|
/ tvm.tirx.Cast(dtype, roi_bin_grid_h),
|
|
roi_start_w
|
|
+ (tvm.tirx.Cast(dtype, rw_idx) + tvm.tirx.const(0.5, dtype))
|
|
* bin_w
|
|
/ tvm.tirx.Cast(dtype, roi_bin_grid_w),
|
|
)
|
|
|
|
if avg_mode:
|
|
return te.sum(
|
|
sample_value(rh, rw) / tvm.tirx.Cast(dtype, count),
|
|
axis=[rh, rw],
|
|
)
|
|
return te.max(sample_value(rh, rw), axis=[rh, rw])
|
|
|
|
|
|
def roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1, aligned=False):
|
|
"""ROI align operator in NCHW layout."""
|
|
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 avg or max. Please pass in a valid mode."
|
|
|
|
_, channel, height, width = data.shape
|
|
num_roi, _ = rois.shape
|
|
dtype = rois.dtype
|
|
|
|
if isinstance(pooled_size, int):
|
|
pooled_size_h = pooled_size_w = pooled_size
|
|
else:
|
|
pooled_size_h, pooled_size_w = pooled_size
|
|
|
|
height_f = tvm.tirx.Cast(dtype, height)
|
|
width_f = tvm.tirx.Cast(dtype, width)
|
|
zero = tvm.tirx.const(0.0, data.dtype)
|
|
|
|
def _bilinear(n, c, y, x):
|
|
outside = tvm.tirx.any(y < -1.0, x < -1.0, y > height_f, x > width_f)
|
|
y = te.min(te.max(y, 0.0), tvm.tirx.Cast(dtype, height - 1))
|
|
x = te.min(te.max(x, 0.0), tvm.tirx.Cast(dtype, width - 1))
|
|
val = bilinear_sample_nchw(data, (n, c, y, x), height - 1, width - 1)
|
|
return tvm.tirx.if_then_else(outside, zero, val)
|
|
|
|
return te.compute(
|
|
(num_roi, channel, pooled_size_h, pooled_size_w),
|
|
lambda i, c, ph, pw: _sample_common(
|
|
i,
|
|
c,
|
|
ph,
|
|
pw,
|
|
rois,
|
|
pooled_size_h,
|
|
pooled_size_w,
|
|
spatial_scale,
|
|
sample_ratio,
|
|
aligned,
|
|
dtype,
|
|
avg_mode,
|
|
_bilinear,
|
|
),
|
|
tag="pool,roi_align_nchw",
|
|
)
|
|
|
|
|
|
def roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1, aligned=False):
|
|
"""ROI align operator in NHWC layout."""
|
|
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 avg or max. Please pass in a valid mode."
|
|
|
|
_, height, width, channel = data.shape
|
|
num_roi, _ = rois.shape
|
|
dtype = rois.dtype
|
|
|
|
if isinstance(pooled_size, int):
|
|
pooled_size_h = pooled_size_w = pooled_size
|
|
else:
|
|
pooled_size_h, pooled_size_w = pooled_size
|
|
|
|
height_f = tvm.tirx.Cast(dtype, height)
|
|
width_f = tvm.tirx.Cast(dtype, width)
|
|
zero = tvm.tirx.const(0.0, data.dtype)
|
|
|
|
def _bilinear(n, c, y, x):
|
|
outside = tvm.tirx.any(y < -1.0, x < -1.0, y > height_f, x > width_f)
|
|
y = te.min(te.max(y, 0.0), tvm.tirx.Cast(dtype, height - 1))
|
|
x = te.min(te.max(x, 0.0), tvm.tirx.Cast(dtype, width - 1))
|
|
val = bilinear_sample_nhwc(data, (n, y, x, c), height - 1, width - 1)
|
|
return tvm.tirx.if_then_else(outside, zero, val)
|
|
|
|
return te.compute(
|
|
(num_roi, pooled_size_h, pooled_size_w, channel),
|
|
lambda i, ph, pw, c: _sample_common(
|
|
i,
|
|
c,
|
|
ph,
|
|
pw,
|
|
rois,
|
|
pooled_size_h,
|
|
pooled_size_w,
|
|
spatial_scale,
|
|
sample_ratio,
|
|
aligned,
|
|
dtype,
|
|
avg_mode,
|
|
_bilinear,
|
|
),
|
|
tag="pool,roi_align_nhwc",
|
|
)
|
|
|
|
|
|
def roi_align(
|
|
data,
|
|
rois,
|
|
pooled_size,
|
|
spatial_scale,
|
|
mode="avg",
|
|
sample_ratio=-1,
|
|
aligned=False,
|
|
layout="NCHW",
|
|
):
|
|
"""ROI align operator."""
|
|
if layout == "NCHW":
|
|
return roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio, aligned)
|
|
if layout == "NHWC":
|
|
return roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio, aligned)
|
|
raise ValueError(f"Unsupported layout for roi_align: {layout}")
|