forked from apache/tvm
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[TOPI, Relay] ROI Pool operator (apache#2811)
- Loading branch information
Showing
13 changed files
with
351 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
# 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]) | ||
roi_start_w = int(round(roi[1] * spatial_scale)) | ||
roi_start_h = int(round(roi[2] * spatial_scale)) | ||
roi_end_w = int(round(roi[3] * spatial_scale)) | ||
roi_end_h = int(round(roi[4] * spatial_scale)) | ||
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 = int(math.floor(ph * bin_h)) | ||
wstart = int(math.floor(pw * bin_w)) | ||
hend = int(math.ceil((ph + 1) * bin_h)) | ||
wend = int(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. | ||
else: | ||
b_np[i, c, ph, pw] = np.max(a_np[batch_index, c, hstart:hend, wstart:wend]) | ||
return b_np |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,5 @@ | ||
# pylint: disable=wildcard-import | ||
"""Faster R-CNN and Mask R-CNN operators""" | ||
from .roi_align import * | ||
from .roi_pool import * | ||
from .proposal import * |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
# pylint: disable=invalid-name | ||
"""ROI pool operator""" | ||
import tvm | ||
from ...util import get_const_tuple | ||
|
||
@tvm.target.generic_func | ||
def roi_pool_nchw(data, rois, pooled_size, spatial_scale): | ||
"""ROI pool operator in NCHW layout. | ||
Parameters | ||
---------- | ||
data : tvm.Tensor | ||
4-D with shape [batch, channel, height, width] | ||
rois : tvm.Tensor | ||
2-D with shape [num_roi, 5]. The last dimension should be in format of | ||
[batch_index, w_start, h_start, w_end, h_end] | ||
pooled_size : int or list/tuple of two ints | ||
output size, or [out_height, out_width] | ||
spatial_scale : float | ||
Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal | ||
of total stride in convolutional layers, which should be in range (0.0, 1.0] | ||
Returns | ||
------- | ||
output : tvm.Tensor | ||
4-D with shape [num_roi, channel, pooled_size, pooled_size] | ||
""" | ||
dtype = rois.dtype | ||
_, channel, height, width = get_const_tuple(data.shape) | ||
num_roi, _ = get_const_tuple(rois.shape) | ||
|
||
if isinstance(pooled_size, int): | ||
pooled_size_h = pooled_size_w = pooled_size | ||
else: | ||
pooled_size_h, pooled_size_w = pooled_size | ||
|
||
def _pool(i, c, ph, pw): | ||
roi = rois[i] | ||
batch_index = roi[0].astype('int32') | ||
roi_start_w, roi_start_h, roi_end_w, roi_end_h = roi[1], roi[2], roi[3], roi[4] | ||
|
||
roi_start_h = tvm.round(roi_start_h * spatial_scale).astype('int32') | ||
roi_start_w = tvm.round(roi_start_w * spatial_scale).astype('int32') | ||
roi_end_h = tvm.round(roi_end_h * spatial_scale).astype('int32') | ||
roi_end_w = tvm.round(roi_end_w * spatial_scale).astype('int32') | ||
|
||
# force malformed ROIs to be 1x1 | ||
roi_h = tvm.max(roi_end_h - roi_start_h + 1, tvm.const(1, 'int32')) | ||
roi_w = tvm.max(roi_end_w - roi_start_w + 1, tvm.const(1, 'int32')) | ||
|
||
bin_h = roi_h.astype(dtype) / pooled_size_h | ||
bin_w = roi_w.astype(dtype) / pooled_size_w | ||
|
||
# use epsilon to prevent floating point precision loss in floor/ceil | ||
epsilon = tvm.const(0.00001, dtype) | ||
hstart = tvm.floor(ph * bin_h + epsilon).astype('int32') | ||
wstart = tvm.floor(pw * bin_w + epsilon).astype('int32') | ||
hend = tvm.ceil((ph + 1) * bin_h - epsilon).astype('int32') | ||
wend = tvm.ceil((pw + 1) * bin_w - epsilon).astype('int32') | ||
hstart = tvm.min(tvm.max(hstart + roi_start_h, 0), height) | ||
wstart = tvm.min(tvm.max(wstart + roi_start_w, 0), width) | ||
hend = tvm.min(tvm.max(hend + roi_start_h, 0), height) | ||
wend = tvm.min(tvm.max(wend + roi_start_w, 0), width) | ||
|
||
non_empty = tvm.all(hstart < hend, wstart < wend) | ||
min_value = lambda dtype: tvm.if_then_else(non_empty, tvm.min_value(dtype), | ||
tvm.const(0.0, dtype)) | ||
# pylint: disable=unnecessary-lambda | ||
_max = tvm.comm_reducer(lambda x, y: tvm.make._OpMax(x, y), min_value, name='max') | ||
rh = tvm.reduce_axis((0, hend - hstart), 'rh') | ||
rw = tvm.reduce_axis((0, wend - wstart), 'rw') | ||
return _max(data[batch_index, c, hstart+rh, wstart+rw], axis=[rh, rw]) | ||
|
||
return tvm.compute((num_roi, channel, pooled_size_h, pooled_size_w), _pool, tag="pool,roi_pool") |
Oops, something went wrong.