diff --git a/mediapipe/calculators/tensorflow/BUILD b/mediapipe/calculators/tensorflow/BUILD index 5f5f51657..c5436bfdb 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 ec7cd70fa..cd0e68a65 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 c4bc819e5..8bd3bf441 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}); diff --git a/mediapipe/framework/BUILD b/mediapipe/framework/BUILD index b289fc582..ca8b0057f 100644 --- a/mediapipe/framework/BUILD +++ b/mediapipe/framework/BUILD @@ -895,6 +895,7 @@ cc_library( "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/log:absl_log", "@com_google_absl//absl/memory", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", ], @@ -996,14 +997,15 @@ cc_library( name = "port", hdrs = ["port.h"], defines = select({ - "//conditions:default": [], - }) + select({ - "//conditions:default": [], - "//mediapipe/gpu:disable_gpu": ["MEDIAPIPE_DISABLE_GPU=1"], - }) + select({ - "//conditions:default": [], - "//mediapipe/framework/port:disable_opencv": ["MEDIAPIPE_DISABLE_OPENCV=1"], - }) + select({ + "//conditions:default": [], + }) + select({ + "//conditions:default": [], + "//mediapipe/gpu:disable_gpu": ["MEDIAPIPE_DISABLE_GPU=1"], + }) + + select({ + "//conditions:default": [], + "//mediapipe/framework/port:disable_opencv": ["MEDIAPIPE_DISABLE_OPENCV=1"], + }) + select({ "//conditions:default": [], # TODO: Improve this. This only sets MEDIAPIPE_DISABLE_OPENCV as a "defines" Make # value, not as a bazel "--define" variable, which has effects in C++ code but not in diff --git a/mediapipe/framework/packet.h b/mediapipe/framework/packet.h index 770dd9d4c..84ae94dc5 100644 --- a/mediapipe/framework/packet.h +++ b/mediapipe/framework/packet.h @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -27,6 +28,7 @@ #include "absl/log/absl_check.h" #include "absl/log/absl_log.h" #include "absl/memory/memory.h" +#include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/synchronization/mutex.h" #include "mediapipe/framework/deps/no_destructor.h" @@ -112,7 +114,18 @@ class Packet { // Transfers the ownership of holder's data to a unique pointer // of the object if the packet is the sole owner of a non-foreign // holder. Otherwise, returns error when the packet can't be consumed. - // See ConsumeOrCopy for threading requirements and example usage. + // + // --- WARNING --- + // Packet is thread-compatible and this member function is non-const. Hence, + // calling it requires exclusive access to the object - callers are + // responsible for ensuring that no other thread is doing anything with the + // packet. + // + // For example, if a node/calculator calls this function, then no other + // calculator should be processing the same packet. Nodes/calculators cannot + // enforce/guarantee this as they don't know of each other, which means graph + // must be written in a special way to account for that. It's error-prone and + // general recommendation is to avoid calling this function. template absl::StatusOr> Consume(); @@ -120,13 +133,22 @@ class Packet { // unique pointer if the packet is the sole owner of a non-foreign // holder. Otherwise, the unique pointer holds a copy of the original // data. In either case, the original packet is set to empty. The - // method returns error when the packet can't be consumed or copied. If + // function returns error when the packet can't be consumed or copied. If // was_copied is not nullptr, it is set to indicate whether the packet // data was copied. - // Packet is thread-compatible, therefore Packet::ConsumeOrCopy() - // must be thread-compatible: clients who use this function are - // responsible for ensuring that no other thread is doing anything - // with the Packet. + // + // --- WARNING --- + // Packet is thread-compatible and this member function is non-const. Hence, + // calling it requires exclusive access to the object - callers are + // responsible for ensuring that no other thread is doing anything with the + // packet. + // + // For example, if a node/calculator calls this function, then no other + // calculator should be processing the same packet. Nodes/calculators cannot + // enforce/guarantee this as they don't know of each other, which means graph + // must be written in a special way to account for that. It's error-prone and + // general recommendation is to avoid calling this function. + // // Example usage: // ASSIGN_OR_RETURN(std::unique_ptr detection, // p.ConsumeOrCopy()); diff --git a/mediapipe/tasks/cc/core/external_file_handler.h b/mediapipe/tasks/cc/core/external_file_handler.h index 3150fde59..01352df7f 100644 --- a/mediapipe/tasks/cc/core/external_file_handler.h +++ b/mediapipe/tasks/cc/core/external_file_handler.h @@ -82,9 +82,11 @@ class ExternalFileHandler { // The aligned mapped memory buffer offset, if any. int64 buffer_aligned_offset_{}; +#ifndef _WIN32 // The aligned mapped memory buffer size in bytes taking into account the // offset shift introduced by buffer_aligned_memory_offset_, if any. int64 buffer_aligned_size_{}; +#endif }; } // namespace core