mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Support string type for split op (#1014)
* Initial commit * Support string type in split * More changes * More changes * Minor nit * Revert nit fix
This commit is contained in:
Hariharan Seshadri
GitHub
parent
31cbb5d33d
commit
49a2928a59
2 changed files with 148 additions and 50 deletions
|
|
@ -16,9 +16,8 @@ ONNX_CPU_OPERATOR_KERNEL(
|
|||
KernelDefBuilder().TypeConstraint("T",
|
||||
std::vector<MLDataType>{
|
||||
DataTypeImpl::GetTensorType<float>(),
|
||||
DataTypeImpl::GetTensorType<double>(),
|
||||
DataTypeImpl::GetTensorType<int32_t>(),
|
||||
}),
|
||||
DataTypeImpl::GetTensorType<std::string>()}),
|
||||
Split);
|
||||
|
||||
Status SplitBase::PrepareForCompute(const TensorShape& input_shape,
|
||||
|
|
@ -36,8 +35,8 @@ Status SplitBase::PrepareForCompute(const TensorShape& input_shape,
|
|||
before_dims = gsl::narrow<int>(input_shape.SizeToDimension(axis));
|
||||
after_dims_including_split_axis = gsl::narrow<int>(input_shape.SizeFromDimension(axis));
|
||||
after_dims_excluding_split = (axis + 1 == num_dimensions)
|
||||
? 1 // we multiply by this value so must be 1 not 0
|
||||
: gsl::narrow<int>(input_shape.SizeFromDimension(axis + 1));
|
||||
? 1 // we multiply by this value so must be 1 not 0
|
||||
: gsl::narrow<int>(input_shape.SizeFromDimension(axis + 1));
|
||||
|
||||
if (split_sizes_.empty()) {
|
||||
// equal split based on number of outputs
|
||||
|
|
@ -73,16 +72,25 @@ Status Split::Compute(OpKernelContext* context) const {
|
|||
status = ComputeImpl<float>(*context, input);
|
||||
else if (data_type == DataTypeImpl::GetType<int32_t>())
|
||||
status = ComputeImpl<int32_t>(*context, input);
|
||||
else if (data_type == DataTypeImpl::GetType<double>()) {
|
||||
/* Need to update CopyMatrix to support double...
|
||||
status = ComputeImpl<double>(*context, input); */
|
||||
ORT_NOT_IMPLEMENTED("Split operator does not support double yet");
|
||||
} else
|
||||
ORT_THROW("Invalid data type for Split operator of ", data_type);
|
||||
else if (data_type == DataTypeImpl::GetType<std::string>())
|
||||
status = ComputeImpl<std::string>(*context, input);
|
||||
else
|
||||
ORT_THROW("Split operator does not support ", data_type, " yet");
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void copy_data(const T* src, T* dst, size_t count) {
|
||||
memcpy(dst, src, count * sizeof(T));
|
||||
}
|
||||
|
||||
template<>
|
||||
inline void copy_data<std::string>(const std::string* src, std::string* dst, size_t count) {
|
||||
const std::string* end = src + count;
|
||||
std::copy(src, end, dst);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status Split::ComputeImpl(OpKernelContext& context, const Tensor& input) const {
|
||||
auto& input_shape = input.Shape();
|
||||
|
|
@ -124,7 +132,7 @@ Status Split::ComputeImpl(OpKernelContext& context, const Tensor& input) const {
|
|||
static_cast<T*>(output_data), // B
|
||||
split_size * after_dims_excluding_split, // ldb
|
||||
[](const T* src, T* dst, size_t count) {
|
||||
memcpy(dst, src, count * sizeof(T));
|
||||
copy_data<T>(src, dst, count);
|
||||
});
|
||||
|
||||
input_offset += split_size * after_dims_excluding_split; // offset by the N data we used in this iteration
|
||||
|
|
|
|||
|
|
@ -7,15 +7,18 @@
|
|||
namespace onnxruntime {
|
||||
namespace test {
|
||||
|
||||
template<class T> using ShapeAndData = std::pair<const std::vector<int64_t>, const std::vector<T>>;
|
||||
template <class T>
|
||||
using ShapeAndData = std::pair<const std::vector<int64_t>, const std::vector<T>>;
|
||||
|
||||
using ShapeAndFloatData = ShapeAndData<float>;
|
||||
using ShapeAndInt32Data = ShapeAndData<int32_t>;
|
||||
using ShapeAndStringData = ShapeAndData<std::string>;
|
||||
using ExpectResult = OpTester::ExpectResult;
|
||||
|
||||
template<typename T> void RunTest(int64_t axis, const std::vector<int64_t> split_sizes, const ShapeAndData<T>& input,
|
||||
const std::vector<ShapeAndData<T>>& outputs, bool is_tensorrt_supported = true,
|
||||
bool expect_failure = false, const std::string& err_msg = {}) {
|
||||
template <typename T>
|
||||
void RunTest(int64_t axis, const std::vector<int64_t> split_sizes, const ShapeAndData<T>& input,
|
||||
const std::vector<ShapeAndData<T>>& outputs, bool is_tensorrt_supported = true,
|
||||
bool expect_failure = false, const std::string& err_msg = {}) {
|
||||
OpTester test("Split");
|
||||
|
||||
test.AddAttribute("axis", axis);
|
||||
|
|
@ -40,16 +43,16 @@ template<typename T> void RunTest(int64_t axis, const std::vector<int64_t> split
|
|||
test.Run(expect_failure ? ExpectResult::kExpectFailure : ExpectResult::kExpectSuccess, err_msg, excluded_providers);
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis0EqualSplit) {
|
||||
TEST(SplitOperatorTest, Axis0EqualSplitFloat) {
|
||||
const int64_t axis = 0;
|
||||
std::vector<ShapeAndFloatData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{4, 2}, // shape
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{1.f, 2.f,
|
||||
|
|
@ -59,7 +62,7 @@ TEST(SplitOperatorTest, Axis0EqualSplit) {
|
|||
{5.f, 6.f,
|
||||
7.f, 8.f}});
|
||||
|
||||
RunTest<float>(axis, {}, input, outputs, false);//TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
RunTest<float>(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis0EqualSplitInt32) {
|
||||
|
|
@ -81,19 +84,41 @@ TEST(SplitOperatorTest, Axis0EqualSplitInt32) {
|
|||
{5, 6,
|
||||
7, 8}});
|
||||
|
||||
RunTest<int32_t>(axis, {}, input, outputs, false);//TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
RunTest<int32_t>(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis0UnequalSplit) {
|
||||
TEST(SplitOperatorTest, Axis0EqualSplitString) {
|
||||
const int64_t axis = 0;
|
||||
std::vector<ShapeAndStringData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndStringData input = {{4, 2}, // shape
|
||||
{"a", "b",
|
||||
"c", "d",
|
||||
"e", "f",
|
||||
"g", "h"}};
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{"a", "b",
|
||||
"c", "d"}});
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{"e", "f",
|
||||
"g", "h"}});
|
||||
|
||||
RunTest<std::string>(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis0UnequalSplitFloat) {
|
||||
const int64_t axis = 0;
|
||||
std::vector<ShapeAndFloatData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{4, 2}, // shape
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
|
||||
std::vector<int64_t> splits{1, 3};
|
||||
|
||||
|
|
@ -104,17 +129,40 @@ TEST(SplitOperatorTest, Axis0UnequalSplit) {
|
|||
5.f, 6.f,
|
||||
7.f, 8.f}});
|
||||
|
||||
RunTest<float>(axis, splits, input, outputs, false);//TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
RunTest<float>(axis, splits, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis1EqualSplit) {
|
||||
TEST(SplitOperatorTest, Axis0UnequalSplitString) {
|
||||
const int64_t axis = 0;
|
||||
std::vector<ShapeAndStringData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndStringData input = {{4, 2}, // shape
|
||||
{"a", "b",
|
||||
"c", "d",
|
||||
"e", "f",
|
||||
"g", "h"}};
|
||||
|
||||
std::vector<int64_t> splits{1, 3};
|
||||
|
||||
outputs.push_back({{1, 2}, {"a", "b"}});
|
||||
|
||||
outputs.push_back({{3, 2},
|
||||
{"c", "d",
|
||||
"e", "f",
|
||||
"g", "h"}});
|
||||
|
||||
RunTest<std::string>(axis, splits, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis1EqualSplitFloat) {
|
||||
const int64_t axis = 1;
|
||||
std::vector<ShapeAndFloatData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{2, 4},
|
||||
{1.f, 2.f, 3.f, 4.f,
|
||||
5.f, 6.f, 7.f, 8.f}};
|
||||
{1.f, 2.f, 3.f, 4.f,
|
||||
5.f, 6.f, 7.f, 8.f}};
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{1.f, 2.f,
|
||||
|
|
@ -127,14 +175,34 @@ TEST(SplitOperatorTest, Axis1EqualSplit) {
|
|||
RunTest<float>(axis, {}, input, outputs, false);
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis1UnequalSplit) {
|
||||
TEST(SplitOperatorTest, Axis1EqualSplitString) {
|
||||
const int64_t axis = 1;
|
||||
std::vector<ShapeAndStringData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndStringData input = {{2, 4},
|
||||
{"a", "b", "c", "d",
|
||||
"e", "f", "g", "h"}};
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{"a", "b",
|
||||
"e", "f"}});
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{"c", "d",
|
||||
"g", "h"}});
|
||||
|
||||
RunTest<std::string>(axis, {}, input, outputs, false);
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis1UnequalSplitFloat) {
|
||||
const int64_t axis = 1;
|
||||
std::vector<ShapeAndFloatData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{2, 4},
|
||||
{1.f, 2.f, 3.f, 4.f,
|
||||
5.f, 6.f, 7.f, 8.f}};
|
||||
{1.f, 2.f, 3.f, 4.f,
|
||||
5.f, 6.f, 7.f, 8.f}};
|
||||
|
||||
std::vector<int64_t> splits{3, 1};
|
||||
|
||||
|
|
@ -149,6 +217,28 @@ TEST(SplitOperatorTest, Axis1UnequalSplit) {
|
|||
RunTest<float>(axis, splits, input, outputs, false);
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, Axis1UnequalSplitString) {
|
||||
const int64_t axis = 1;
|
||||
std::vector<ShapeAndStringData> outputs;
|
||||
|
||||
// input shape and data
|
||||
ShapeAndStringData input = {{2, 4},
|
||||
{"a", "b", "c", "d",
|
||||
"e", "f", "g", "h"}};
|
||||
|
||||
std::vector<int64_t> splits{3, 1};
|
||||
|
||||
outputs.push_back({{2, 3},
|
||||
{"a", "b", "c",
|
||||
"e", "f", "g"}});
|
||||
|
||||
outputs.push_back({{2, 1},
|
||||
{"d",
|
||||
"h"}});
|
||||
|
||||
RunTest<std::string>(axis, splits, input, outputs, false);
|
||||
}
|
||||
|
||||
ShapeAndFloatData CreateInput(std::vector<int64_t> shape) {
|
||||
auto size = TensorShape(shape).Size();
|
||||
|
||||
|
|
@ -280,8 +370,8 @@ TEST(SplitOperatorTest, NegativeAxis) {
|
|||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{2, 4},
|
||||
{1.f, 2.f, 3.f, 4.f,
|
||||
5.f, 6.f, 7.f, 8.f}};
|
||||
{1.f, 2.f, 3.f, 4.f,
|
||||
5.f, 6.f, 7.f, 8.f}};
|
||||
|
||||
outputs.push_back({{2, 2},
|
||||
{1.f, 2.f,
|
||||
|
|
@ -300,10 +390,10 @@ TEST(SplitOperatorTest, InvalidAxis) {
|
|||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{4, 2}, // shape
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
|
||||
outputs.push_back({{1}, {0.f}});
|
||||
|
||||
|
|
@ -317,17 +407,17 @@ TEST(SplitOperatorTest, SplitAttributeSumTooSmall) {
|
|||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{4, 2}, // shape
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
|
||||
std::vector<int64_t> splits{1, 2}; // should sum to 4
|
||||
|
||||
outputs.push_back({{1, 2}, {1.f, 2.f}});
|
||||
outputs.push_back({{2, 2}, {3.f, 4.f, 5.f, 6.f}});
|
||||
|
||||
RunTest<float>(axis, splits, input, outputs, false, true, "Cannot split using values in 'split' attribute");//TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
RunTest<float>(axis, splits, input, outputs, false, true, "Cannot split using values in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
TEST(SplitOperatorTest, InvalidValueInSplitAttribute) {
|
||||
|
|
@ -336,16 +426,16 @@ TEST(SplitOperatorTest, InvalidValueInSplitAttribute) {
|
|||
|
||||
// input shape and data
|
||||
ShapeAndFloatData input = {{4, 2}, // shape
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
{1.f, 2.f,
|
||||
3.f, 4.f,
|
||||
5.f, 6.f,
|
||||
7.f, 8.f}};
|
||||
|
||||
std::vector<int64_t> splits{1, 0, 3}; // 0 is not valid
|
||||
outputs.push_back({{1, 2}, {1.f, 2.f}});
|
||||
outputs.push_back({{3, 2}, {3.f, 4.f, 5.f, 6.f, 7.f, 8.f}});
|
||||
|
||||
RunTest<float>(axis, splits, input, outputs, false, true, "Invalid value in 'split' attribute");//TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
RunTest<float>(axis, splits, input, outputs, false, true, "Invalid value in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM
|
||||
}
|
||||
|
||||
/*
|
||||
|
|
|
|||
Loading…
Reference in a new issue