cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Modifying tensor_to_vector_float_calculator to take in D_BFLOAT16 values

Browse Source 
PiperOrigin-RevId: 565189254
This commit is contained in:
MediaPipe Team 2023-09-13 16:07:58 -07:00 committed by Copybara-Service
parent 6dc1239aa9
commit e1d1877e07
3 changed files with 92 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mediapipe/calculators/tensorflow/BUILD
View File

@ -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",
],
)

31
mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc
View File

@ -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,23 +80,33 @@ 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>();
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());
} else {
if (!options_.flatten_nd()) {
@ -101,11 +115,18 @@ absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) {
<< "tensor shape is: " << input_tensor.shape().DebugString();
}
auto output =
absl::make_unique<std::vector<float>>(input_tensor.NumElements());
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());
}

59
mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc
View File

@ -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});