saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-28 20:11:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update transpose logic for Scan 9 outputs so it's the inverse of the input transpose logic. (#401)

Browse source
This commit is contained in:
Scott McKay 2019-01-29 22:17:04 +10:00 committed by GitHub
parent b194b7df0d
commit d4d3270891
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 87 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/core/providers/cpu/controlflow/scan_9.cc
View file

@ -343,7 +343,7 @@ Status ScanImpl::SetupInputs() {
std::vector<int64_t> permutations;
std::vector<int64_t> new_shape;
CalculateTransposedShape(input_shape, sequence_dim, permutations, new_shape);
CalculateTransposedShapeForInput(input_shape, sequence_dim, permutations, new_shape);
if (!alloc) {
status = context_.GetTempSpaceAllocator(&alloc);
@ -474,7 +474,7 @@ Status ScanImpl::TransposeOutput() {
std::vector<int64_t> permutations;
std::vector<int64_t> new_shape;
CalculateTransposedShape(temporary_output_tensor.Shape(), axis, permutations, new_shape);
CalculateTransposedShapeForOutput(temporary_output_tensor.Shape(), axis, permutations, new_shape);
Tensor* output = context_.Output(output_index, new_shape);
ORT_ENFORCE(output, "Outputs from Scan are not optional and should never be null.");

36
onnxruntime/core/providers/cpu/controlflow/scan_utils.cc
View file

@ -211,25 +211,47 @@ MLValue AllocateTensorInMLValue(const MLDataType data_type, const TensorShape& s
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc()};
};
void CalculateTransposedShape(const TensorShape& input_shape, int64_t axis,
std::vector<int64_t>& permutations, std::vector<int64_t>& output_shape) {
int64_t rank = input_shape.NumDimensions();
const auto& dims = input_shape.GetDims();
void CalculateTransposedShapeForInput(const TensorShape& original_shape, int64_t axis,
std::vector<int64_t>& permutations, std::vector<int64_t>& transposed_shape) {
int64_t rank = original_shape.NumDimensions();
const auto& dims = original_shape.GetDims();
permutations.reserve(rank);
permutations.push_back(axis);
output_shape.reserve(rank);
output_shape.push_back(dims[axis]);
transposed_shape.reserve(rank);
transposed_shape.push_back(dims[axis]);
for (int64_t i = 0; i < rank; ++i) {
if (i != axis) {
permutations.push_back(i);
output_shape.push_back(dims[i]);
transposed_shape.push_back(dims[i]);
}
}
}
void CalculateTransposedShapeForOutput(const TensorShape& original_shape, int64_t axis,
std::vector<int64_t>& permutations, std::vector<int64_t>& transposed_shape) {
int64_t rank = original_shape.NumDimensions();
const auto& dims = original_shape.GetDims();
permutations.reserve(rank);
transposed_shape.reserve(rank);
for (int64_t i = 1; i <= axis; ++i) {
permutations.push_back(i);
transposed_shape.push_back(dims[i]);
}
permutations.push_back(0);
transposed_shape.push_back(dims[0]);
for (int64_t i = axis + 1; i < rank; ++i) {
permutations.push_back(i);
transposed_shape.push_back(dims[i]);
}
}
LoopStateVariable::LoopStateVariable(const MLValue& original_value,
MLValue& final_value,
const int64_t sequence_len,

15
onnxruntime/core/providers/cpu/controlflow/scan_utils.h
View file

@ -173,14 +173,23 @@ Status IterateSequence(OpKernelContextInternal& context,
MLValue AllocateTensorInMLValue(const MLDataType data_type, const TensorShape& shape, AllocatorPtr& allocator);
/**
Calculate the transpose permutations and output shape by shifting the chosen axis to the first dimension.
Calculate the transpose permutations and shape by shifting the chosen axis TO the first dimension.
The other dimension indexes or values are pushed in order after the chosen axis.
e.g. if shape is {2, 3, 4} and axis 1 is chosen the permutations will be {1, 0, 2} and output shape will be {3, 2, 4}
if axis 2 is chosen the permutations will be {2, 0, 1} and the output shape will be {4, 2, 3}
*/
void CalculateTransposedShape(const TensorShape& input_shape, int64_t axis,
std::vector<int64_t>& permutations, std::vector<int64_t>& output_shape);
void CalculateTransposedShapeForInput(const TensorShape& original_shape, int64_t axis,
std::vector<int64_t>& permutations, std::vector<int64_t>& transposed_shape);
/**
Calculate the transpose permutations and shape by shifting the chosen axis FROM the first dimension.
e.g. if shape is {4, 2, 3} and axis 2 is chosen, dimension 0 will move to dimension 2,
the permutations will be {1, 2, 0} and output shape will be {2, 3, 4}
*/
void CalculateTransposedShapeForOutput(const TensorShape& original_shape, int64_t axis,
std::vector<int64_t>& permutations, std::vector<int64_t>& transposed_shape);
} // namespace detail
} // namespace scan

46
onnxruntime/test/providers/cpu/controlflow/scan_test.cc
View file

@ -382,8 +382,8 @@ static void RunTest_v9(const std::string test_name, int64_t sequence_len, int64_
if (axis >= -rank && axis < rank) {
std::vector<int64_t> permutations;
std::vector<int64_t> new_shape;
scan::detail::CalculateTransposedShape(output_shape, HandleNegativeAxis(axis, output_shape.size()),
permutations, new_shape);
scan::detail::CalculateTransposedShapeForOutput(output_shape, HandleNegativeAxis(axis, output_shape.size()),
permutations, new_shape);
return new_shape;
}
}
@ -844,6 +844,48 @@ TEST(Scan9, TransposeOutput) {
iteration_count_out, output_0, output_1, output_2, output_3);
}
TEST(Scan9, TransposeOutputDim2) {
// Construct scan body subgraph with 1 scan inputs, 1 scan outputs
// scan-in-1 => scan-out-1
Model model("ScanBody");
auto& graph = model.MainGraph();
TypeProto float_tensor;
float_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1);
float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1);
auto& scan_in_1 = graph.GetOrCreateNodeArg("scan_in_1", &float_tensor);
auto& scan_out_1 = graph.GetOrCreateNodeArg("scan_out_1", &float_tensor);
graph.AddNode("pass_through", "Identity", "Copy scan_in_1 to scan_out_1", {&scan_in_1}, {&scan_out_1});
auto status = graph.Resolve();
EXPECT_EQ(status, Status::OK());
auto& scan_body = graph.ToGraphProto();
ScanOpTester test{9};
std::vector<int64_t> input_shape{2, 1, 1};
// transpose on axis 2, so dim 0 of the output (copied directly from input of {2, 1, 1})
// will move to dim 2 of the output giving shape {1, 1, 2}
std::vector<int64_t> output_axes{2};
std::vector<int64_t> output_shape{1, 1, 2};
test.AddAttribute("body", scan_body);
test.AddAttribute<int64_t>("num_scan_inputs", 1);
test.AddAttribute<std::vector<int64_t>>("scan_output_axes", output_axes);
// the data won't change, but the shape should be transposed from 2, 1, 1 to 1, 1, 2, which
// OpTester::Run will validate
test.AddInput<float>("scan_input_1", input_shape, {1.0, 2.0});
test.AddOutput<float>("scan_output_1", output_shape, {1.0, 2.0});
test.Run();
}
static void InvalidInput(bool is_v8) {
const int64_t batch_size = 1;
const int64_t sequence_len = 2;