add frontend resize kernel (#194)

* - add frontend resize kernel

* - fix resize test

* - fix bug
- add onnx test for resize

* fix: modify codes as reviewer suggested

---------

Co-authored-by: Haojie Wang <haojie0429@gmail.com>
This commit is contained in:
Chenjie Duan 2023-12-29 13:32:56 +08:00 committed by GitHub
parent 3967b437c8
commit 83f1de93d0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 187 additions and 0 deletions

View File

@ -65,6 +65,12 @@ class GraphHandlerObj {
std::optional<float> max);
Tensor transpose(Tensor data, Tensor transposed, Shape perm);
Tensor reshape(Tensor data, Tensor reshaped, Shape shape);
Tensor resize(Tensor input, Tensor output,
const std::optional<vector<int>> &axes, Tensor sizes,
Tensor scales, Tensor roi, vector<uint32_t> sizes_,
vector<float> scales_, vector<float> roi_, string mode,
string ratioPolicy, string nearestMode,
string coordTransMode);
Tensor concat(TensorVec inputs, Tensor output, int dim);
Tensor attentionKVCache(Tensor input_k_cache, Tensor input_v_cache,
Tensor input_q, Tensor input_k, Tensor input_v,

View File

@ -27,6 +27,60 @@ class ResizeObj : public OperatorObj {
enum class EKeepAspectRatioPolicy { stretch, notLarger, notSmaller, none };
enum class ECoeffMode { nearest, linear, cubic };
static ECoordinateTransMode fromECoordinateTransModeStr(string mode) {
if (mode == "half_pixel") {
return ECoordinateTransMode::halfPixel;
} else if (mode == "asymmetric") {
return ECoordinateTransMode::asymmetric;
} else if (mode == "align_corners") {
return ECoordinateTransMode::alignCorners;
} else if (mode == "pytorch_half_pixel") {
return ECoordinateTransMode::pytorchHalfPixel;
} else if (mode == "tf_crop_and_resize") {
return ECoordinateTransMode::tfCropAndResize;
} else {
IT_TODO_HALT();
}
}
static ENearestMode fromENearestModeStr(string mode) {
if (mode == "round_prefer_floor") {
return ENearestMode::roundPreferFloor;
} else if (mode == "round_prefer_ceil") {
return ENearestMode::roundPreferCeil;
} else if (mode == "floor") {
return ENearestMode::floor;
} else if (mode == "ceil") {
return ENearestMode::ceil;
} else {
return ENearestMode::none;
}
}
static EKeepAspectRatioPolicy fromRatioPolicyStr(string ratioPolicyStr) {
if (ratioPolicyStr == "stretch") {
return EKeepAspectRatioPolicy::stretch;
} else if (ratioPolicyStr == "not_larger") {
return EKeepAspectRatioPolicy::notLarger;
} else if (ratioPolicyStr == "not_smaller") {
return EKeepAspectRatioPolicy::notSmaller;
} else {
return EKeepAspectRatioPolicy::none;
}
}
static ECoeffMode fromECoeffModeStr(string mode) {
if (mode == "nearest") {
return ECoeffMode::nearest;
} else if (mode == "linear") {
return ECoeffMode::linear;
} else if (mode == "cubic") {
return ECoeffMode::cubic;
} else {
IT_TODO_HALT();
}
}
private:
vector<int> axes;
vector<float> scales;

View File

@ -535,6 +535,65 @@ class OnnxStub:
tensors.get(node.output[0]),
shape,
)
elif node.op_type == "Resize":
output = tensors.get(node.output[0])
attributes = _parse_attribute(
node,
{
"antialias": 0,
"axes": None,
"coordinate_transformation_mode": "half_pixel",
"cubic_coeff_a": -0.75,
"exclude_outside": 0,
"extrapolation_value": 0.0,
"keep_aspect_ratio_policy": "none",
"mode": "nearest",
"nearest_mode": "none",
},
)
(
axes,
keep_aspect_ratio_policy,
coordinate_transformation_mode,
mode,
nearest_mode,
) = (
attributes[name]
for name in [
"axes",
"keep_aspect_ratio_policy",
"coordinate_transformation_mode",
"mode",
"nearest_mode",
]
)
if len(node.input) > 1:
roiVal = _parse_data(data[node.input[1]])
else:
roiVal = []
if len(node.input) > 2:
scalesVal = _parse_data(data[node.input[2]])
else:
scalesVal = []
if len(node.input) > 3:
sizesVal = _parse_data(data[node.input[3]])
else:
sizesVal = []
tensors[node.output[0]] = self.handler.resize(
tensors[node.input[0]],
output,
axes,
tensors[node.input[3]] if len(node.input) > 3 else None,
tensors[node.input[2]] if len(node.input) > 2 else None,
tensors[node.input[1]] if len(node.input) > 1 else None,
sizesVal,
scalesVal,
roiVal,
mode,
keep_aspect_ratio_policy,
nearest_mode,
coordinate_transformation_mode,
)
elif node.op_type == "Squeeze":
input_shape = _search_shape(model, node.input[0])
axes = set(

View File

@ -295,6 +295,14 @@ class TestStringMethods(unittest.TestCase):
make_graph([reshape], "reshape", [data, shape], [reshaped], [shape_data])
)
def test_resize(self):
x = make_tensor_value_info("x", TensorProto.FLOAT, [1, 128, 40, 40])
roi = make_tensor("roi", TensorProto.FLOAT, [0], [])
scales = make_tensor("scales", TensorProto.FLOAT, [4], [1, 1, 2, 2])
y = make_tensor_value_info("y", TensorProto.FLOAT, [1, 128, 80, 80])
reshape = make_node("Resize", ["x", "roi", "scales"], ["y"], name="resize")
make_and_import_model(make_graph([reshape], "resize", [x], [y], [roi, scales]))
def test_concat(self):
input1 = make_tensor_value_info("input1", TensorProto.FLOAT, [1, 3, 2, 4])
input2 = make_tensor_value_info("input2", TensorProto.FLOAT, [1, 3, 2, 5])

View File

@ -17,6 +17,7 @@
#include "operators/recv.h"
#include "operators/reduce.h"
#include "operators/reshape.h"
#include "operators/resize.h"
#include "operators/send.h"
#include "operators/slice.h"
#include "operators/softmax.h"
@ -254,6 +255,64 @@ Tensor GraphHandlerObj::reshape(Tensor data, Tensor reshaped, Shape shape) {
}
}
Tensor GraphHandlerObj::resize(Tensor input, Tensor output,
const std::optional<vector<int>> &axes,
Tensor sizes, Tensor scales, Tensor roi,
vector<uint32_t> sizes_, vector<float> scales_,
vector<float> roi_, string mode,
string ratioPolicy, string nearestMode,
string coordTransMode) {
if (sizes_.size() > 0) {
sizes->dataMalloc();
sizes->copyin<uint32_t>(sizes_);
}
if (scales_.size() > 0) {
scales->dataMalloc();
scales->copyin<float>(scales_);
}
if (roi_.size() > 0) {
roi->dataMalloc();
roi->copyin<float>(roi_);
}
ResizeObj::EKeepAspectRatioPolicy ratioPolicy_ =
ResizeObj::fromRatioPolicyStr(ratioPolicy);
ResizeObj::ENearestMode nearestMode_ =
ResizeObj::fromENearestModeStr(nearestMode);
ResizeObj::ECoordinateTransMode coordTransMode_ =
ResizeObj::fromECoordinateTransModeStr(coordTransMode);
ResizeObj::ECoeffMode mode_ = ResizeObj::fromECoeffModeStr(mode);
if (output) {
if (mode == "nearest") {
g->addOpWithOutputs<ResizeObj>(
std::move(input), output, std::move(axes), std::move(sizes),
std::move(scales), std::move(roi), ratioPolicy_, nearestMode_,
coordTransMode_);
} else {
g->addOpWithOutputs<ResizeObj>(
std::move(input), output, std::move(axes), std::move(sizes),
std::move(scales), std::move(roi), mode_, ratioPolicy_,
coordTransMode_);
}
return output;
} else {
if (mode == "nearest") {
return g
->addOp<ResizeObj>(std::move(input), output, std::move(axes),
std::move(sizes), std::move(scales),
std::move(roi), ratioPolicy_, nearestMode_,
coordTransMode_)
->getOutput();
} else {
return g
->addOp<ResizeObj>(std::move(input), output, std::move(axes),
std::move(sizes), std::move(scales),
std::move(roi), mode_, ratioPolicy_,
coordTransMode_)
->getOutput();
}
}
}
Tensor GraphHandlerObj::concat(TensorVec inputs, Tensor output, int dim) {
if (output) {
g->addOpWithOutputs<ConcatObj>(std::move(inputs), output, dim);

View File

@ -506,6 +506,7 @@ void init_graph_builder(py::module &m) {
.def("transpose", &Handler::transpose, policy::move)
.def("depthToSpace", &Handler::depthToSpace, policy::move)
.def("reshape", &Handler::reshape, policy::move)
.def("resize", &Handler::resize, policy::move)
.def("concat", &Handler::concat, policy::move)
.def("attentionKVCache", &Handler::attentionKVCache, policy::move)
.def("split", &Handler::split, policy::move)