Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
Xiaomi
Mace
提交
433d74a7
Mace
项目概览
Xiaomi
/
Mace
通知
106
Star
40
Fork
27
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
Mace
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
433d74a7
编写于
3月 22, 2018
作者:
L
liuqi
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Add PSROIAlign op.
上级
cd859dc9
变更
6
显示空白变更内容
内联
并排
Showing
6 changed file
with
307 addition
and
52 deletion
+307
-52
mace/core/operator.cc
mace/core/operator.cc
+2
-0
mace/kernels/proposal.h
mace/kernels/proposal.h
+65
-52
mace/kernels/psroi_align.h
mace/kernels/psroi_align.h
+175
-0
mace/ops/proposal_test.cc
mace/ops/proposal_test.cc
+2
-0
mace/ops/psroi_align.cc
mace/ops/psroi_align.cc
+19
-0
mace/ops/psroi_align.h
mace/ops/psroi_align.h
+44
-0
未找到文件。
mace/core/operator.cc
浏览文件 @
433d74a7
...
...
@@ -90,6 +90,7 @@ extern void Register_Eltwise(OperatorRegistry *op_registry);
extern
void
Register_FullyConnected
(
OperatorRegistry
*
op_registry
);
extern
void
Register_Slice
(
OperatorRegistry
*
op_registry
);
extern
void
Register_Proposal
(
OperatorRegistry
*
op_registry
);
extern
void
Register_PSROIAlign
(
OperatorRegistry
*
op_registry
);
}
// namespace ops
...
...
@@ -120,6 +121,7 @@ OperatorRegistry::OperatorRegistry() {
ops
::
Register_FullyConnected
(
this
);
ops
::
Register_Slice
(
this
);
ops
::
Register_Proposal
(
this
);
ops
::
Register_PSROIAlign
(
this
);
}
}
// namespace mace
mace/kernels/proposal.h
浏览文件 @
433d74a7
...
...
@@ -42,33 +42,33 @@ std::vector<std::vector<float>> GenerateAnchors(
std
::
vector
<
std
::
vector
<
float
>>
anchors
(
scales_size
*
ratios_size
,
std
::
vector
<
float
>
(
4
));
int
idx
=
0
;
std
::
vector
<
float
>
tmp_anchor
(
4
);
#pragma omp parallel for
for
(
size_t
ratio_idx
=
0
;
ratio_idx
<
ratios_size
;
++
ratio_idx
)
{
float
ws
=
::
roundf
(
::
sqrtf
(
size
/
ratios
[
ratio_idx
]));
float
hs
=
::
roundf
(
ws
*
ratios
[
ratio_idx
]);
std
::
vector
<
float
>
tmp_anchor
(
4
);
tmp_anchor
[
0
]
=
base_window
[
2
]
-
(
ws
-
1
)
/
2
;
tmp_anchor
[
1
]
=
base_window
[
3
]
-
(
hs
-
1
)
/
2
;
tmp_anchor
[
2
]
=
base_window
[
2
]
+
(
ws
-
1
)
/
2
;
tmp_anchor
[
3
]
=
base_window
[
3
]
+
(
hs
-
1
)
/
2
;
auto
window
=
WHCenters
(
tmp_anchor
);
for
(
size_t
scale_idx
=
0
;
scale_idx
<
scales_size
;
++
scale_idx
)
{
const
size_t
idx
=
ratio_idx
*
scales_size
+
scale_idx
;
ws
=
window
[
0
]
*
scales
[
scale_idx
];
hs
=
window
[
1
]
*
scales
[
scale_idx
];
anchors
[
idx
][
0
]
=
window
[
2
]
-
(
ws
-
1
)
/
2
;
anchors
[
idx
][
1
]
=
window
[
3
]
-
(
hs
-
1
)
/
2
;
anchors
[
idx
][
2
]
=
window
[
2
]
+
(
ws
-
1
)
/
2
;
anchors
[
idx
][
3
]
=
window
[
3
]
+
(
hs
-
1
)
/
2
;
idx
++
;
}
}
return
anchors
;
}
std
::
vector
<
int
>
nms
(
const
float
*
bboxes_ptr
,
const
float
*
scores_ptr
,
const
index_t
num_bboxes
,
const
float
thresh
)
{
const
float
thresh
,
const
int
post_nms_top_n
)
{
std
::
vector
<
int
>
keep
;
std
::
vector
<
int
>
suppressed
(
num_bboxes
,
0
);
...
...
@@ -79,33 +79,29 @@ std::vector<int> nms(const float *bboxes_ptr,
(
bboxes_ptr
[
idx
+
3
]
-
bboxes_ptr
[
idx
+
1
]
+
1
);
}
std
::
vector
<
int
>
order
(
num_bboxes
,
0
);
iota
(
order
.
begin
(),
order
.
end
(),
0
);
for
(
int
i
=
0
;
i
<
num_bboxes
;
++
i
)
{
if
(
suppressed
[
i
]
==
1
)
continue
;
keep
.
push_back
(
i
);
if
(
keep
.
size
()
>=
post_nms_top_n
)
break
;
int
coord_idx
=
i
<<
2
;
const
float
x1
=
bboxes_ptr
[
coord_idx
];
const
float
y1
=
bboxes_ptr
[
coord_idx
+
1
];
const
float
x2
=
bboxes_ptr
[
coord_idx
+
2
];
const
float
y2
=
bboxes_ptr
[
coord_idx
+
3
];
const
float
area1
=
areas
[
i
];
for
(
int
j
=
i
+
1
;
j
<
num_bboxes
;
++
j
)
{
if
(
suppressed
[
j
]
==
1
)
continue
;
float
x1
,
y1
,
x2
,
y2
,
area1
,
iou
;
for
(
index_t
i
=
0
;
i
<
num_bboxes
;
++
i
)
{
int
idx
=
order
[
i
];
if
(
suppressed
[
idx
]
==
1
)
continue
;
keep
.
push_back
(
idx
);
int
coord_idx
=
idx
<<
2
;
x1
=
bboxes_ptr
[
coord_idx
];
y1
=
bboxes_ptr
[
coord_idx
+
1
];
x2
=
bboxes_ptr
[
coord_idx
+
2
];
y2
=
bboxes_ptr
[
coord_idx
+
3
];
area1
=
areas
[
idx
];
for
(
index_t
j
=
i
+
1
;
j
<
num_bboxes
;
++
j
)
{
const
int
other_idx
=
order
[
j
];
if
(
suppressed
[
other_idx
]
==
1
)
continue
;
coord_idx
=
other_idx
<<
2
;
iou
=
std
::
max
<
float
>
(
0.0
,
coord_idx
=
j
<<
2
;
const
float
iou
=
std
::
max
<
float
>
(
0.0
,
std
::
min
(
x2
,
bboxes_ptr
[
coord_idx
+
2
])
-
std
::
max
(
x1
,
bboxes_ptr
[
coord_idx
])
+
1
)
*
std
::
max
<
float
>
(
0.0
,
std
::
min
(
y2
,
bboxes_ptr
[
coord_idx
+
3
])
-
std
::
max
(
y1
,
bboxes_ptr
[
coord_idx
+
1
])
+
1
);
if
((
iou
/
(
area1
+
areas
[
other_idx
]
-
iou
))
>=
thresh
)
{
suppressed
[
other_idx
]
=
1
;
if
((
iou
/
(
area1
+
areas
[
j
]
-
iou
))
>=
thresh
)
{
suppressed
[
j
]
=
1
;
}
}
}
...
...
@@ -140,8 +136,8 @@ struct ProposalFunctor {
MACE_CHECK
((
rpn_cls_prob
->
dim
(
3
)
/
2
==
rpn_bbox_pred
->
dim
(
3
)
/
4
)
&&
(
rpn_cls_prob
->
dim
(
3
)
/
2
==
anchors_
.
size
()));
const
float
*
img_info
=
img_info_tensor
->
data
<
float
>
();
const
in
dex_t
im_height
=
img_info
[
0
]
-
1
;
const
in
dex_t
im_width
=
img_info
[
1
]
-
1
;
const
in
t
im_height
=
static_cast
<
int
>
(
img_info
[
0
]
-
1
)
;
const
in
t
im_width
=
static_cast
<
int
>
(
img_info
[
1
]
-
1
)
;
const
index_t
feat_height
=
rpn_cls_prob
->
dim
(
1
);
const
index_t
feat_width
=
rpn_cls_prob
->
dim
(
2
);
const
int
anchors_size
=
anchors_
.
size
();
...
...
@@ -150,35 +146,35 @@ struct ProposalFunctor {
std
::
vector
<
std
::
vector
<
float
>>
proposals
(
anchors_size
*
feat_height
*
feat_width
,
std
::
vector
<
float
>
(
4
));
int
shift_w
,
shift_h
;
int
sanc_idx
=
0
;
#pragma omp parallel for collapse(3)
for
(
int
h_idx
=
0
;
h_idx
<
feat_height
;
++
h_idx
)
{
shift_h
=
h_idx
*
feat_stride_
;
for
(
int
w_idx
=
0
;
w_idx
<
feat_width
;
++
w_idx
)
{
shift_w
=
w_idx
*
feat_stride_
;
for
(
int
a_idx
=
0
;
a_idx
<
anchors_size
;
++
a_idx
)
{
const
int
shift_h
=
h_idx
*
feat_stride_
;
const
int
shift_w
=
w_idx
*
feat_stride_
;
const
index_t
sanc_idx
=
(
h_idx
*
feat_width
+
w_idx
)
*
anchors_size
+
a_idx
;
proposals
[
sanc_idx
][
0
]
=
anchors_
[
a_idx
][
0
]
+
shift_w
;
proposals
[
sanc_idx
][
1
]
=
anchors_
[
a_idx
][
1
]
+
shift_h
;
proposals
[
sanc_idx
][
2
]
=
anchors_
[
a_idx
][
2
]
+
shift_w
;
proposals
[
sanc_idx
][
3
]
=
anchors_
[
a_idx
][
3
]
+
shift_h
;
sanc_idx
++
;
}
}
}
// Convert anchors into proposals via bbox transformations
// 2. clip predicted boxes to image
std
::
vector
<
int
>
keep
;
const
index_t
min_size
=
min_size_
*
img_info
[
2
];
// 3. remove predicted boxes with either height or width < threshold
// (NOTE: convert min_size to input image scale stored in im_info[2])
const
float
*
bbox_deltas
=
rpn_bbox_pred
->
data
<
float
>
();
sanc_idx
=
0
;
#pragma omp parallel for collapse(3)
for
(
int
h_idx
=
0
;
h_idx
<
feat_height
;
++
h_idx
)
{
for
(
int
w_idx
=
0
;
w_idx
<
feat_width
;
++
w_idx
)
{
for
(
int
a_idx
=
0
;
a_idx
<
anchors_size
;
++
a_idx
)
{
float
width
=
proposals
[
sanc_idx
][
2
]
-
proposals
[
sanc_idx
][
0
]
+
1
;
float
height
=
proposals
[
sanc_idx
][
3
]
-
proposals
[
sanc_idx
][
1
]
+
1
;
const
int
sanc_idx
=
(
h_idx
*
feat_width
+
w_idx
)
*
anchors_size
+
a_idx
;
const
float
width
=
proposals
[
sanc_idx
][
2
]
-
proposals
[
sanc_idx
][
0
]
+
1
;
const
float
height
=
proposals
[
sanc_idx
][
3
]
-
proposals
[
sanc_idx
][
1
]
+
1
;
int
delta_offset
=
sanc_idx
*
4
;
float
pred_ctr_x
=
bbox_deltas
[
delta_offset
+
0
]
*
width
+
(
proposals
[
sanc_idx
][
0
]
+
width
/
2
);
...
...
@@ -200,12 +196,25 @@ struct ProposalFunctor {
std
::
min
<
float
>
(
pred_ctr_y
+
pred_h
/
2
,
im_height
),
0
);
width
=
proposals
[
sanc_idx
][
2
]
-
proposals
[
sanc_idx
][
0
]
+
1
;
height
=
proposals
[
sanc_idx
][
3
]
-
proposals
[
sanc_idx
][
1
]
+
1
;
}
}
}
// 3. remove predicted boxes with either height or width < threshold
// (NOTE: convert min_size to input image scale stored in im_info[2])
std
::
vector
<
int
>
keep
;
const
float
min_size
=
min_size_
*
img_info
[
2
];
for
(
int
h_idx
=
0
;
h_idx
<
feat_height
;
++
h_idx
)
{
for
(
int
w_idx
=
0
;
w_idx
<
feat_width
;
++
w_idx
)
{
for
(
int
a_idx
=
0
;
a_idx
<
anchors_size
;
++
a_idx
)
{
const
int
sanc_idx
=
(
h_idx
*
feat_width
+
w_idx
)
*
anchors_size
+
a_idx
;
const
float
width
=
proposals
[
sanc_idx
][
2
]
-
proposals
[
sanc_idx
][
0
]
+
1
;
const
float
height
=
proposals
[
sanc_idx
][
3
]
-
proposals
[
sanc_idx
][
1
]
+
1
;
if
(
width
>=
min_size
&&
height
>=
min_size
)
{
keep
.
push_back
(
sanc_idx
);
}
sanc_idx
++
;
}
}
}
...
...
@@ -227,6 +236,7 @@ struct ProposalFunctor {
int
size
=
std
::
min
<
int
>
(
pre_nms_top_n_
,
keep
.
size
());
std
::
vector
<
float
>
nms_scores
(
size
,
0
);
std
::
vector
<
float
>
nms_proposals
((
size
<<
2
),
0
);
#pragma omp parallel for
for
(
int
i
=
0
;
i
<
size
;
++
i
)
{
nms_scores
[
i
]
=
scores
[
score_idx_func
(
keep
[
i
])];
nms_proposals
[
i
<<
2
]
=
proposals
[
keep
[
i
]][
0
];
...
...
@@ -238,15 +248,18 @@ struct ProposalFunctor {
/* 6. apply nms (e.g. threshold = 0.7)
7. take after_nms_topN (e.g. 300)
8. return the top proposals (-> RoIs top) */
auto
nms_result
=
nms
(
nms_proposals
.
data
(),
nms_scores
.
data
(),
nms_scores
.
size
(),
thresh_
);
auto
nms_result
=
nms
(
nms_proposals
.
data
(),
nms_scores
.
size
(),
thresh_
,
post_nms_top_n_
);
// Output rois blob
// Our RPN implementation only supports a single input image, so all
// batch inds are 0
size
=
st
d
::
min
<
int
>
(
post_nms_top_n_
,
nms_result
.
size
());
output
->
Resize
({
1
,
1
,
size
,
5
});
size
=
st
atic_cast
<
int
>
(
nms_result
.
size
());
output
->
Resize
({
size
,
1
,
1
,
5
});
auto
output_ptr
=
output
->
mutable_data
<
float
>
();
#pragma omp parallel for
for
(
int
i
=
0
;
i
<
size
;
++
i
)
{
const
int
out_idx
=
i
*
5
;
const
int
nms_idx
=
nms_result
[
i
]
*
4
;
...
...
mace/kernels/psroi_align.h
0 → 100644
浏览文件 @
433d74a7
//
// Copyright (c) 2017 XiaoMi All rights reserved.
//
#ifndef MACE_KERNELS_PSROI_ALIGN_H_
#define MACE_KERNELS_PSROI_ALIGN_H_
#include "mace/core/future.h"
#include "mace/core/tensor.h"
#include "mace/public/mace.h"
namespace
mace
{
namespace
kernels
{
template
<
DeviceType
D
,
typename
T
>
struct
PSROIAlignFunctor
{
PSROIAlignFunctor
(
const
T
spatial_scale
,
const
int
output_dim
,
const
int
group_size
)
:
spatial_scale_
(
spatial_scale
),
output_dim_
(
output_dim
),
group_size_
(
group_size
)
{}
void
operator
()(
const
Tensor
*
input
,
const
Tensor
*
rois
,
Tensor
*
output
,
StatsFuture
*
future
)
{
const
int
height
=
static_cast
<
int
>
(
input
->
dim
(
1
));
const
int
width
=
static_cast
<
int
>
(
input
->
dim
(
2
));
const
int
channels
=
static_cast
<
int
>
(
input
->
dim
(
3
));
const
int
pooled_height
=
group_size_
;
const
int
pooled_width
=
group_size_
;
const
T
*
input_ptr
=
input
->
data
<
T
>
();
const
T
*
rois_ptr
=
rois
->
data
<
T
>
();
// Number of ROIs
const
int
num_rois
=
rois
->
dim
(
0
);
const
int
batch_size
=
input
->
dim
(
0
);
output
->
Resize
({
num_rois
,
pooled_height
,
pooled_width
,
output_dim_
});
T
*
output_ptr
=
output
->
mutable_data
<
T
>
();
for
(
int
n
=
0
;
n
<
num_rois
;
++
n
)
{
int
roi_batch_ind
=
rois_ptr
[
0
];
T
roi_start_w
=
static_cast
<
T
>
(
rois_ptr
[
1
])
*
spatial_scale_
;
T
roi_start_h
=
static_cast
<
T
>
(
rois_ptr
[
2
])
*
spatial_scale_
;
T
roi_end_w
=
static_cast
<
T
>
(
rois_ptr
[
3
]
+
1.
)
*
spatial_scale_
;
T
roi_end_h
=
static_cast
<
T
>
(
rois_ptr
[
4
]
+
1.
)
*
spatial_scale_
;
MACE_CHECK
(
roi_batch_ind
>=
0
);
MACE_CHECK
(
roi_batch_ind
<
batch_size
);
// Force too small ROIs to be 1x1
T
roi_width
=
std
::
max
(
roi_end_w
-
roi_start_w
,
static_cast
<
T
>
(
0.1
));
T
roi_height
=
std
::
max
(
roi_end_h
-
roi_start_h
,
static_cast
<
T
>
(
0.1
));
// Compute w and h at bottom
T
bin_size_h
=
roi_height
/
static_cast
<
T
>
(
pooled_height
);
T
bin_size_w
=
roi_width
/
static_cast
<
T
>
(
pooled_width
);
const
T
*
batch_data
=
input_ptr
+
roi_batch_ind
*
height
*
width
*
channels
;
std
::
vector
<
T
>
vhstart
,
vwstart
,
vhend
,
vwend
;
for
(
int
ph
=
0
;
ph
<
pooled_height
;
++
ph
)
{
for
(
int
pw
=
0
;
pw
<
pooled_width
;
++
pw
)
{
T
hstart
=
static_cast
<
T
>
(
ph
)
*
bin_size_h
+
roi_start_h
;
T
wstart
=
static_cast
<
T
>
(
pw
)
*
bin_size_w
+
roi_start_w
;
T
hend
=
static_cast
<
T
>
(
ph
+
1
)
*
bin_size_h
+
roi_start_h
;
T
wend
=
static_cast
<
T
>
(
pw
+
1
)
*
bin_size_w
+
roi_start_w
;
// Add roi offsets and clip to input boundaries
hstart
=
std
::
min
(
std
::
max
(
hstart
,
static_cast
<
T
>
(
0.
)),
static_cast
<
T
>
(
height
));
hend
=
std
::
min
(
std
::
max
(
hend
,
static_cast
<
T
>
(
0.
)),
static_cast
<
T
>
(
height
));
wstart
=
std
::
min
(
std
::
max
(
wstart
,
static_cast
<
T
>
(
0.
)),
static_cast
<
T
>
(
width
));
wend
=
std
::
min
(
std
::
max
(
wend
,
static_cast
<
T
>
(
0.
)),
static_cast
<
T
>
(
width
));
vhstart
.
push_back
(
hstart
);
vwstart
.
push_back
(
wstart
);
vhend
.
push_back
(
hend
);
vwend
.
push_back
(
wend
);
}
}
#pragma omp parallel for collapse(3)
for
(
int
ph
=
0
;
ph
<
pooled_height
;
++
ph
)
{
for
(
int
pw
=
0
;
pw
<
pooled_width
;
++
pw
)
{
for
(
int
c
=
0
;
c
<
output_dim_
;
++
c
)
{
const
int
pool_index
=
ph
*
pooled_width
+
pw
;
const
int
out_idx
=
pool_index
*
output_dim_
+
c
;
const
int
in_chan_idx
=
(
c
*
pooled_height
+
ph
)
*
pooled_width
+
pw
;
T
hstart
=
vhstart
[
pool_index
];
T
hend
=
vhend
[
pool_index
];
T
wstart
=
vwstart
[
pool_index
];
T
wend
=
vwend
[
pool_index
];
bool
is_empty
=
(
hend
<=
hstart
)
||
(
wend
<=
wstart
);
T
out_sum
=
0
;
for
(
T
h
=
hstart
;
h
<
hend
;
h
+=
1.
)
{
for
(
T
w
=
wstart
;
w
<
wend
;
w
+=
1.
)
{
// Selecting four regular locations for bilinear interpolation
int
x_left
=
std
::
floor
(
w
);
int
x_right
=
std
::
ceil
(
w
);
int
y_bottom
=
std
::
floor
(
h
);
int
y_top
=
std
::
ceil
(
h
);
int
top_left_index
=
(
y_top
*
width
+
x_left
)
*
channels
+
in_chan_idx
;
int
top_right_index
=
(
y_top
*
width
+
x_right
)
*
channels
+
in_chan_idx
;
int
bottom_left_index
=
(
y_bottom
*
width
+
x_left
)
*
channels
+
in_chan_idx
;
int
bottom_right_index
=
(
y_bottom
*
width
+
x_right
)
*
channels
+
in_chan_idx
;
bool
is_top_left_in
=
x_left
>=
0
&&
x_left
<=
width
-
1
&&
y_top
>=
0
&&
y_top
<=
height
-
1
;
bool
is_top_right_in
=
x_right
>=
0
&&
x_right
<=
width
-
1
&&
y_top
>=
0
&&
y_top
<=
height
-
1
;
bool
is_bottom_left_in
=
x_left
>=
0
&&
x_left
<=
width
-
1
&&
y_bottom
>=
0
&&
y_bottom
<=
height
-
1
;
bool
is_bottom_right_in
=
x_right
>=
0
&&
x_right
<=
width
-
1
&&
y_bottom
>=
0
&&
y_bottom
<=
height
-
1
;
if
(
is_top_left_in
)
{
out_sum
+=
(
1
-
w
+
x_left
)
*
(
1
-
y_top
+
h
)
*
batch_data
[
top_left_index
];
}
if
(
is_top_right_in
)
{
out_sum
+=
(
1
-
x_right
+
w
)
*
(
1
-
y_top
+
h
)
*
batch_data
[
top_right_index
];
}
if
(
is_bottom_left_in
)
{
out_sum
+=
(
1
-
w
+
x_left
)
*
(
1
-
h
+
y_bottom
)
*
batch_data
[
bottom_left_index
];
}
if
(
is_bottom_right_in
)
{
out_sum
+=
(
1
-
x_right
+
w
)
*
(
1
-
h
+
y_bottom
)
*
batch_data
[
bottom_right_index
];
}
}
}
T
bin_area
=
(
hend
-
hstart
)
*
(
wend
-
wstart
);
output_ptr
[
out_idx
]
=
is_empty
?
0.
:
out_sum
/
bin_area
;
}
}
}
// Increment ROI data pointer
rois_ptr
+=
5
;
output_ptr
+=
pooled_height
*
pooled_width
*
output_dim_
;
}
}
const
T
spatial_scale_
;
const
int
output_dim_
;
const
int
group_size_
;
};
}
// namepsace kernels
}
// namespace mace
#endif // MACE_KERNELS_PSROI_ALIGN_H_
mace/ops/proposal_test.cc
浏览文件 @
433d74a7
...
...
@@ -4,6 +4,7 @@
#include "mace/core/operator.h"
#include "mace/ops/ops_test_util.h"
#include <fstream>
namespace
mace
{
namespace
ops
{
...
...
@@ -53,6 +54,7 @@ TEST_F(ProposalOpTest, CPUSimple) {
auto
expected_tensor
=
CreateTensor
<
float
>
({
1
,
1
,
1
,
5
},
{
0
,
0
,
0
,
255
,
255
});
ExpectTensorNear
<
float
>
(
*
expected_tensor
,
*
net
.
GetTensor
(
"Output"
),
1e-5
);
}
...
...
mace/ops/psroi_align.cc
0 → 100644
浏览文件 @
433d74a7
//
// Copyright (c) 2017 XiaoMi All rights reserved.
//
#include "mace/ops/psroi_align.h"
namespace
mace
{
namespace
ops
{
void
Register_PSROIAlign
(
OperatorRegistry
*
op_registry
)
{
REGISTER_OPERATOR
(
op_registry
,
OpKeyBuilder
(
"PSROIAlign"
)
.
Device
(
DeviceType
::
CPU
)
.
TypeConstraint
<
float
>
(
"T"
)
.
Build
(),
PSROIAlignOp
<
DeviceType
::
CPU
,
float
>
);
}
}
// namespace ops
}
// namespace mace
mace/ops/psroi_align.h
0 → 100644
浏览文件 @
433d74a7
//
// Copyright (c) 2017 XiaoMi All rights reserved.
//
#ifndef MACE_OPS_PSROI_ALIGN_H_
#define MACE_OPS_PSROI_ALIGN_H_
#include "mace/core/operator.h"
#include "mace/kernels/psroi_align.h"
namespace
mace
{
namespace
ops
{
template
<
DeviceType
D
,
class
T
>
class
PSROIAlignOp
:
public
Operator
<
D
,
T
>
{
public:
PSROIAlignOp
(
const
OperatorDef
&
operator_def
,
Workspace
*
ws
)
:
Operator
<
D
,
T
>
(
operator_def
,
ws
),
functor_
(
OperatorBase
::
GetSingleArgument
<
T
>
(
"spatial_scale"
,
0
),
OperatorBase
::
GetSingleArgument
<
int
>
(
"output_dim"
,
0
),
OperatorBase
::
GetSingleArgument
<
int
>
(
"group_size"
,
0
))
{}
bool
Run
(
StatsFuture
*
future
)
override
{
const
Tensor
*
input
=
this
->
Input
(
INPUT
);
const
Tensor
*
rois
=
this
->
Input
(
ROIS
);
Tensor
*
output
=
this
->
Output
(
OUTPUT
);
functor_
(
input
,
rois
,
output
,
future
);
return
true
;
}
private:
kernels
::
PSROIAlignFunctor
<
D
,
T
>
functor_
;
protected:
OP_INPUT_TAGS
(
INPUT
,
ROIS
);
OP_OUTPUT_TAGS
(
OUTPUT
);
};
}
// namespace ops
}
// namespace mace
#endif // MACE_OPS_PSROI_ALIGN_H_
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录