Modifying tensor_to_vector_float_calculator to take in D_BFLOAT16 values
PiperOrigin-RevId: 565189254
This commit is contained in:
parent
6dc1239aa9
commit
e1d1877e07
|
@ -723,6 +723,7 @@ cc_library(
|
|||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core/platform:bfloat16",
|
||||
] + select({
|
||||
"//conditions:default": [
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
|
@ -1139,6 +1140,7 @@ cc_test(
|
|||
"//mediapipe/util:packet_test_util",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core/platform:bfloat16",
|
||||
],
|
||||
)
|
||||
|
||||
|
|
|
@ -15,12 +15,16 @@
|
|||
// Calculator converts from one-dimensional Tensor of DT_FLOAT to vector<float>
|
||||
// OR from (batched) two-dimensional Tensor of DT_FLOAT to vector<vector<float>.
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_options.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "tensorflow/core/framework/tensor.h"
|
||||
#include "tensorflow/core/framework/types.h"
|
||||
#include "tensorflow/core/platform/bfloat16.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
|
@ -76,21 +80,31 @@ absl::Status TensorToVectorFloatCalculator::Open(CalculatorContext* cc) {
|
|||
absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) {
|
||||
const tf::Tensor& input_tensor =
|
||||
cc->Inputs().Index(0).Value().Get<tf::Tensor>();
|
||||
RET_CHECK(tf::DT_FLOAT == input_tensor.dtype())
|
||||
<< "expected DT_FLOAT input but got "
|
||||
RET_CHECK(tf::DT_FLOAT == input_tensor.dtype() ||
|
||||
tf::DT_BFLOAT16 == input_tensor.dtype())
|
||||
<< "expected DT_FLOAT or DT_BFLOAT_16 input but got "
|
||||
<< tensorflow::DataTypeString(input_tensor.dtype());
|
||||
|
||||
if (options_.tensor_is_2d()) {
|
||||
RET_CHECK(2 == input_tensor.dims())
|
||||
<< "Expected 2-dimensional Tensor, but the tensor shape is: "
|
||||
<< input_tensor.shape().DebugString();
|
||||
auto output = absl::make_unique<std::vector<std::vector<float>>>(
|
||||
auto output = std::make_unique<std::vector<std::vector<float>>>(
|
||||
input_tensor.dim_size(0), std::vector<float>(input_tensor.dim_size(1)));
|
||||
for (int i = 0; i < input_tensor.dim_size(0); ++i) {
|
||||
auto& instance_output = output->at(i);
|
||||
const auto& slice = input_tensor.Slice(i, i + 1).unaligned_flat<float>();
|
||||
for (int j = 0; j < input_tensor.dim_size(1); ++j) {
|
||||
instance_output.at(j) = slice(j);
|
||||
if (tf::DT_BFLOAT16 == input_tensor.dtype()) {
|
||||
const auto& slice =
|
||||
input_tensor.Slice(i, i + 1).unaligned_flat<tf::bfloat16>();
|
||||
for (int j = 0; j < input_tensor.dim_size(1); ++j) {
|
||||
instance_output.at(j) = static_cast<float>(slice(j));
|
||||
}
|
||||
} else {
|
||||
const auto& slice =
|
||||
input_tensor.Slice(i, i + 1).unaligned_flat<float>();
|
||||
for (int j = 0; j < input_tensor.dim_size(1); ++j) {
|
||||
instance_output.at(j) = slice(j);
|
||||
}
|
||||
}
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
|
@ -101,10 +115,17 @@ absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) {
|
|||
<< "tensor shape is: " << input_tensor.shape().DebugString();
|
||||
}
|
||||
auto output =
|
||||
absl::make_unique<std::vector<float>>(input_tensor.NumElements());
|
||||
const auto& tensor_values = input_tensor.unaligned_flat<float>();
|
||||
for (int i = 0; i < input_tensor.NumElements(); ++i) {
|
||||
output->at(i) = tensor_values(i);
|
||||
std::make_unique<std::vector<float>>(input_tensor.NumElements());
|
||||
if (tf::DT_BFLOAT16 == input_tensor.dtype()) {
|
||||
const auto& tensor_values = input_tensor.unaligned_flat<tf::bfloat16>();
|
||||
for (int i = 0; i < input_tensor.NumElements(); ++i) {
|
||||
output->at(i) = static_cast<float>(tensor_values(i));
|
||||
}
|
||||
} else {
|
||||
const auto& tensor_values = input_tensor.unaligned_flat<float>();
|
||||
for (int i = 0; i < input_tensor.NumElements(); ++i) {
|
||||
output->at(i) = tensor_values(i);
|
||||
}
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
|
|
@ -12,6 +12,8 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_options.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
|
@ -19,6 +21,7 @@
|
|||
#include "mediapipe/util/packet_test_util.h"
|
||||
#include "tensorflow/core/framework/tensor.h"
|
||||
#include "tensorflow/core/framework/types.pb.h"
|
||||
#include "tensorflow/core/platform/bfloat16.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
|
@ -72,6 +75,62 @@ TEST_F(TensorToVectorFloatCalculatorTest, ConvertsToVectorFloat) {
|
|||
}
|
||||
}
|
||||
|
||||
TEST_F(TensorToVectorFloatCalculatorTest, CheckBFloat16Type) {
|
||||
SetUpRunner(false, false);
|
||||
const tf::TensorShape tensor_shape(std::vector<tf::int64>{5});
|
||||
auto tensor = std::make_unique<tf::Tensor>(tf::DT_BFLOAT16, tensor_shape);
|
||||
auto tensor_vec = tensor->vec<tf::bfloat16>();
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
tensor_vec(i) = static_cast<tf::bfloat16>(1 << i);
|
||||
}
|
||||
|
||||
const int64_t time = 1234;
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(tensor.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const std::vector<float>& output_vector =
|
||||
output_packets[0].Get<std::vector<float>>();
|
||||
|
||||
EXPECT_EQ(5, output_vector.size());
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
const float expected = static_cast<float>(1 << i);
|
||||
EXPECT_EQ(expected, output_vector[i]);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TensorToVectorFloatCalculatorTest, CheckBFloat16TypeAllDim) {
|
||||
SetUpRunner(false, true);
|
||||
const tf::TensorShape tensor_shape(std::vector<tf::int64>{2, 2, 2});
|
||||
auto tensor = std::make_unique<tf::Tensor>(tf::DT_BFLOAT16, tensor_shape);
|
||||
auto slice = tensor->flat<tf::bfloat16>();
|
||||
for (int i = 0; i < 2 * 2 * 2; ++i) {
|
||||
// 2^i can be represented exactly in floating point numbers if 'i' is small.
|
||||
slice(i) = static_cast<tf::bfloat16>(1 << i);
|
||||
}
|
||||
|
||||
const int64_t time = 1234;
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(tensor.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const std::vector<float>& output_vector =
|
||||
output_packets[0].Get<std::vector<float>>();
|
||||
EXPECT_EQ(2 * 2 * 2, output_vector.size());
|
||||
for (int i = 0; i < 2 * 2 * 2; ++i) {
|
||||
const float expected = static_cast<float>(1 << i);
|
||||
EXPECT_EQ(expected, output_vector[i]);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TensorToVectorFloatCalculatorTest, ConvertsBatchedToVectorVectorFloat) {
|
||||
SetUpRunner(true, false);
|
||||
const tf::TensorShape tensor_shape(std::vector<tf::int64>{1, 5});
|
||||
|
|
Loading…
Reference in New Issue
Block a user