diff --git a/mediapipe/calculators/tensorflow/BUILD b/mediapipe/calculators/tensorflow/BUILD index 5f5f5165..c5436bfd 100644 --- a/mediapipe/calculators/tensorflow/BUILD +++ b/mediapipe/calculators/tensorflow/BUILD @@ -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", ], ) diff --git a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc index ec7cd70f..cd0e68a6 100644 --- a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc +++ b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc @@ -15,12 +15,16 @@ // Calculator converts from one-dimensional Tensor of DT_FLOAT to vector // OR from (batched) two-dimensional Tensor of DT_FLOAT to vector. +#include +#include + #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(); - 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>>( + auto output = std::make_unique>>( input_tensor.dim_size(0), std::vector(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(); - 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(); + for (int j = 0; j < input_tensor.dim_size(1); ++j) { + instance_output.at(j) = static_cast(slice(j)); + } + } else { + const auto& slice = + input_tensor.Slice(i, i + 1).unaligned_flat(); + 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>(input_tensor.NumElements()); - const auto& tensor_values = input_tensor.unaligned_flat(); - for (int i = 0; i < input_tensor.NumElements(); ++i) { - output->at(i) = tensor_values(i); + std::make_unique>(input_tensor.NumElements()); + if (tf::DT_BFLOAT16 == input_tensor.dtype()) { + const auto& tensor_values = input_tensor.unaligned_flat(); + for (int i = 0; i < input_tensor.NumElements(); ++i) { + output->at(i) = static_cast(tensor_values(i)); + } + } else { + const auto& tensor_values = input_tensor.unaligned_flat(); + for (int i = 0; i < input_tensor.NumElements(); ++i) { + output->at(i) = tensor_values(i); + } } cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp()); } diff --git a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc index c4bc819e..8bd3bf44 100644 --- a/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc +++ b/mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include + #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{5}); + auto tensor = std::make_unique(tf::DT_BFLOAT16, tensor_shape); + auto tensor_vec = tensor->vec(); + for (int i = 0; i < 5; ++i) { + tensor_vec(i) = static_cast(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& output_packets = + runner_->Outputs().Index(0).packets; + EXPECT_EQ(1, output_packets.size()); + EXPECT_EQ(time, output_packets[0].Timestamp().Value()); + const std::vector& output_vector = + output_packets[0].Get>(); + + EXPECT_EQ(5, output_vector.size()); + for (int i = 0; i < 5; ++i) { + const float expected = static_cast(1 << i); + EXPECT_EQ(expected, output_vector[i]); + } +} + +TEST_F(TensorToVectorFloatCalculatorTest, CheckBFloat16TypeAllDim) { + SetUpRunner(false, true); + const tf::TensorShape tensor_shape(std::vector{2, 2, 2}); + auto tensor = std::make_unique(tf::DT_BFLOAT16, tensor_shape); + auto slice = tensor->flat(); + 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(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& output_packets = + runner_->Outputs().Index(0).packets; + EXPECT_EQ(1, output_packets.size()); + EXPECT_EQ(time, output_packets[0].Timestamp().Value()); + const std::vector& output_vector = + output_packets[0].Get>(); + EXPECT_EQ(2 * 2 * 2, output_vector.size()); + for (int i = 0; i < 2 * 2 * 2; ++i) { + const float expected = static_cast(1 << i); + EXPECT_EQ(expected, output_vector[i]); + } +} + TEST_F(TensorToVectorFloatCalculatorTest, ConvertsBatchedToVectorVectorFloat) { SetUpRunner(true, false); const tf::TensorShape tensor_shape(std::vector{1, 5});