Fix bug in format converter.

PiperOrigin-RevId: 297738172
Change-Id: I8150a9a9712b4f2118a924699bddff72d1e339f9
This commit is contained in:
Yunlu Li 2020-02-27 18:19:32 -08:00 committed by TensorFlower Gardener
parent a962580295
commit d63a01cb1f
2 changed files with 64 additions and 2 deletions

View File

@ -285,10 +285,12 @@ void FormatConverter<T>::Populate(const T* src_data, std::vector<int> indices,
}
const int metadata_idx = 2 * level;
const int shape_of_level = dim_metadata_[metadata_idx][0];
if (format_[level] == kTfLiteDimDense) {
for (int i = 0; i < dim_metadata_[metadata_idx][0]; i++) {
for (int i = 0; i < shape_of_level; i++) {
indices[level] = i;
Populate(src_data, indices, level + 1, i, src_data_ptr);
Populate(src_data, indices, level + 1, prev_idx * shape_of_level + i,
src_data_ptr);
}
} else {
const auto& array_segments = dim_metadata_[metadata_idx];

View File

@ -230,6 +230,66 @@ TEST(FormatConverterTest, SimpleTestS1S0) {
EXPECT_EQ(data_back, dense_values);
}
TEST(FormatConverterTest, 3DTestS0D1S2) {
const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
const std::vector<int> dense_shape = {3, 2, 2};
const std::vector<int> traversal_order = {0, 1, 2};
const std::vector<TfLiteDimensionType> format = {
kTfLiteDimSparseCSR, kTfLiteDimDense, kTfLiteDimSparseCSR};
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
const std::vector<int> dm0_0 = {0, 2};
const std::vector<int> dm0_1 = {0, 2};
const std::vector<int> dm1 = {2};
const std::vector<int> dm2_0 = {0, 1, 3, 4, 5};
const std::vector<int> dm2_1 = {0, 0, 1, 0, 1};
EXPECT_EQ(dm0_0, dim_metadata[0]);
EXPECT_EQ(dm0_1, dim_metadata[1]);
EXPECT_EQ(dm1, dim_metadata[2]);
EXPECT_EQ(dm2_0, dim_metadata[4]);
EXPECT_EQ(dm2_1, dim_metadata[5]);
const auto& data = converter.GetData();
const std::vector<int> expected_data = {6, 9, 8, 5, 7};
EXPECT_EQ(expected_data, data);
converter.SparseToDense(expected_data.data());
const auto& data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
TEST(FormatConverterTest, 3DTestD0D1S2) {
const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
const std::vector<int> dense_shape = {3, 2, 2};
const std::vector<int> traversal_order = {0, 1, 2};
const std::vector<TfLiteDimensionType> format = {
kTfLiteDimDense, kTfLiteDimDense, kTfLiteDimSparseCSR};
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
const std::vector<int> dm0 = {3};
const std::vector<int> dm1 = {2};
const std::vector<int> dm2_0 = {0, 1, 3, 3, 3, 4, 5};
const std::vector<int> dm2_1 = {0, 0, 1, 0, 1};
EXPECT_EQ(dm0, dim_metadata[0]);
EXPECT_EQ(dm1, dim_metadata[2]);
EXPECT_EQ(dm2_0, dim_metadata[4]);
EXPECT_EQ(dm2_1, dim_metadata[5]);
const auto& data = converter.GetData();
const std::vector<int> expected_data = {6, 9, 8, 5, 7};
EXPECT_EQ(expected_data, data);
converter.SparseToDense(expected_data.data());
const auto& data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
TEST(FormatConverterTest, 3DTestS0S1S2) {
const std::vector<int> dense_values = {1, 7, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 2, 0, 0, 4, 8, 3, 9};