forked from jiuyuan/InfiniTensor
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:
parent
3967b437c8
commit
83f1de93d0
|
@ -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,
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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])
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue