Make TensorToVectorFloatCalculator compatible with unaligned tensors.
No performance impact is expected, since the unaligned Eigen::TensorMap is used only to populate a std::vector<float>. PiperOrigin-RevId: 505251810
This commit is contained in:
parent
702cc0c42c
commit
ee2f940e1f
|
@ -1054,6 +1054,7 @@ cc_test(
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/util:packet_test_util",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
|
|
|
@ -102,7 +102,7 @@ absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) {
|
||||||
}
|
}
|
||||||
auto output =
|
auto output =
|
||||||
absl::make_unique<std::vector<float>>(input_tensor.NumElements());
|
absl::make_unique<std::vector<float>>(input_tensor.NumElements());
|
||||||
const auto& tensor_values = input_tensor.flat<float>();
|
const auto& tensor_values = input_tensor.unaligned_flat<float>();
|
||||||
for (int i = 0; i < input_tensor.NumElements(); ++i) {
|
for (int i = 0; i < input_tensor.NumElements(); ++i) {
|
||||||
output->at(i) = tensor_values(i);
|
output->at(i) = tensor_values(i);
|
||||||
}
|
}
|
||||||
|
|
|
@ -16,6 +16,7 @@
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/calculator_runner.h"
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
#include "mediapipe/framework/port/gtest.h"
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/util/packet_test_util.h"
|
||||||
#include "tensorflow/core/framework/tensor.h"
|
#include "tensorflow/core/framework/tensor.h"
|
||||||
#include "tensorflow/core/framework/types.pb.h"
|
#include "tensorflow/core/framework/types.pb.h"
|
||||||
|
|
||||||
|
@ -129,5 +130,28 @@ TEST_F(TensorToVectorFloatCalculatorTest, FlattenShouldTakeAllDimensions) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(TensorToVectorFloatCalculatorTest, AcceptsUnalignedTensors) {
|
||||||
|
SetUpRunner(/*tensor_is_2d=*/false, /*flatten_nd=*/false);
|
||||||
|
|
||||||
|
const tf::TensorShape tensor_shape(std::vector<tf::int64>{2, 5});
|
||||||
|
tf::Tensor tensor(tf::DT_FLOAT, tensor_shape);
|
||||||
|
auto slice = tensor.Slice(1, 1).flat<float>();
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
slice(i) = i;
|
||||||
|
}
|
||||||
|
|
||||||
|
auto input_tensor = tensor.SubSlice(1);
|
||||||
|
// Ensure that the input tensor is unaligned.
|
||||||
|
ASSERT_FALSE(input_tensor.IsAligned());
|
||||||
|
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||||
|
MakePacket<tf::Tensor>(input_tensor).At(Timestamp(5)));
|
||||||
|
|
||||||
|
ASSERT_TRUE(runner_->Run().ok());
|
||||||
|
|
||||||
|
EXPECT_THAT(runner_->Outputs().Index(0).packets,
|
||||||
|
ElementsAre(PacketContainsTimestampAndPayload<std::vector<float>>(
|
||||||
|
Timestamp(5), std::vector<float>({0, 1, 2, 3, 4}))));
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|
Loading…
Reference in New Issue
Block a user