diff --git a/tensorflow/lite/experimental/resource/BUILD b/tensorflow/lite/experimental/resource/BUILD index b1b53afa45b..de4ff592716 100644 --- a/tensorflow/lite/experimental/resource/BUILD +++ b/tensorflow/lite/experimental/resource/BUILD @@ -26,3 +26,17 @@ cc_library( "//tensorflow/lite/kernels/internal:tensor", ], ) + +cc_test( + name = "resource_variable_test", + srcs = [ + "resource_variable_test.cc", + ], + deps = [ + ":resource", + "//tensorflow/lite:util", + "//tensorflow/lite/c:common", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tensorflow/lite/experimental/resource/resource_variable.cc b/tensorflow/lite/experimental/resource/resource_variable.cc index c16db39047c..eb170757aa0 100644 --- a/tensorflow/lite/experimental/resource/resource_variable.cc +++ b/tensorflow/lite/experimental/resource/resource_variable.cc @@ -69,6 +69,8 @@ TfLiteStatus ResourceVariable::AssignFrom(const TfLiteTensor* tensor) { tensor_.data.raw = old_raw; if (old_bytes != tensor->bytes) { TfLiteTensorRealloc(tensor->bytes, &tensor_); + } else { + tensor_.bytes = old_bytes; } memcpy(tensor_.data.raw, tensor->data.raw, tensor_.bytes); @@ -82,8 +84,8 @@ void CreateResourceVariableIfNotAvailable(ResourceMap* resources, if (resources->count(resource_id) != 0) { return; } - resources->emplace( - resource_id, std::unique_ptr(new ResourceVariable())); + resources->emplace(resource_id, + std::unique_ptr(new ResourceVariable())); } ResourceVariable* GetResourceVariable(ResourceMap* resources, int resource_id) { diff --git a/tensorflow/lite/experimental/resource/resource_variable_test.cc b/tensorflow/lite/experimental/resource/resource_variable_test.cc new file mode 100644 index 00000000000..92a1e8a29d4 --- /dev/null +++ b/tensorflow/lite/experimental/resource/resource_variable_test.cc @@ -0,0 +1,173 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include "tensorflow/lite/experimental/resource/resource_variable.h" + +#include + +#include +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/util.h" + +namespace tflite { +namespace resource { +// Helper util that initialize 'tensor'. +void InitTensor(const std::vector& shape, TfLiteAllocationType alloc_type, + float default_value, TfLiteTensor* tensor) { + memset(tensor, 0, sizeof(TfLiteTensor)); + int num_elements = 1; + for (auto dim : shape) num_elements *= dim; + if (shape.empty()) num_elements = 0; + float* buf = static_cast(malloc(sizeof(float) * num_elements)); + for (int i = 0; i < num_elements; ++i) buf[i] = default_value; + const int bytes = num_elements * sizeof(buf[0]); + auto* dims = ConvertArrayToTfLiteIntArray(shape.size(), shape.data()); + TfLiteTensorReset(TfLiteType::kTfLiteFloat32, nullptr, dims, {}, + reinterpret_cast(buf), bytes, alloc_type, nullptr, + false, tensor); +} + +TEST(ResourceTest, NonDynamicTensorAssign) { + ResourceVariable var; + EXPECT_FALSE(var.IsInitialized()); + + TfLiteTensor tensor; + std::vector shape = {1}; + InitTensor(shape, kTfLiteArenaRw, 1.0f, &tensor); + + EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor)); + EXPECT_TRUE(var.IsInitialized()); + auto* value = var.GetTensor(); + + // Variables are always dynamic type. + EXPECT_EQ(kTfLiteDynamic, value->allocation_type); + EXPECT_EQ(kTfLiteFloat32, value->type); + EXPECT_EQ(sizeof(float), value->bytes); + EXPECT_EQ(1, value->dims->size); + EXPECT_EQ(1, value->dims->data[0]); + EXPECT_EQ(1.0f, value->data.f[0]); + + // Cleanup + // For non dynamic tensors we need to delete the buffers manually. + free(tensor.data.raw); + TfLiteTensorFree(&tensor); +} + +TEST(ResourceTest, DynamicTensorAssign) { + ResourceVariable var; + EXPECT_FALSE(var.IsInitialized()); + + TfLiteTensor tensor; + std::vector shape = {1}; + InitTensor(shape, kTfLiteDynamic, 1.0f, &tensor); + + EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor)); + EXPECT_TRUE(var.IsInitialized()); + auto* value = var.GetTensor(); + + // Variables are always dynamic type. + EXPECT_EQ(kTfLiteDynamic, value->allocation_type); + EXPECT_EQ(kTfLiteFloat32, value->type); + EXPECT_EQ(sizeof(float), value->bytes); + EXPECT_EQ(1, value->dims->size); + EXPECT_EQ(1, value->dims->data[0]); + EXPECT_EQ(1.0f, value->data.f[0]); + + // Cleanup + TfLiteTensorFree(&tensor); +} + +TEST(ResourceTest, AssignSameSizeTensor) { + ResourceVariable var; + EXPECT_FALSE(var.IsInitialized()); + + // We create 2 tensors and make 2 calls for Assign. + // The second Assign call should trigger the case of assign with same size. + TfLiteTensor tensor_a, tensor_b; + std::vector shape_a = {1}; + std::vector shape_b = {1}; + InitTensor(shape_a, kTfLiteDynamic, 1.0, &tensor_a); + InitTensor(shape_b, kTfLiteDynamic, 4.0, &tensor_b); + + EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_a)); + EXPECT_TRUE(var.IsInitialized()); + auto* value = var.GetTensor(); + // Variables are always dynamic type. + EXPECT_EQ(kTfLiteDynamic, value->allocation_type); + EXPECT_EQ(kTfLiteFloat32, value->type); + EXPECT_EQ(sizeof(float), value->bytes); + EXPECT_EQ(1, value->dims->size); + EXPECT_EQ(1, value->dims->data[0]); + EXPECT_EQ(1.0f, value->data.f[0]); + + // Second AssignFrom but now tensor_b has same size as the variable. + EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_b)); + EXPECT_TRUE(var.IsInitialized()); + value = var.GetTensor(); + // Variables are always dynamic type. + EXPECT_EQ(kTfLiteDynamic, value->allocation_type); + EXPECT_EQ(kTfLiteFloat32, value->type); + EXPECT_EQ(sizeof(float), value->bytes); + EXPECT_EQ(1, value->dims->size); + EXPECT_EQ(1, value->dims->data[0]); + EXPECT_EQ(4.0f, value->data.f[0]); + + // Cleanup + TfLiteTensorFree(&tensor_a); + TfLiteTensorFree(&tensor_b); +} + +TEST(ResourceTest, AssignDifferentSizeTensor) { + ResourceVariable var; + EXPECT_FALSE(var.IsInitialized()); + + // We create 2 tensors and make 2 calls for Assign. + // The second Assign call should trigger the case of assign with different + // size. + TfLiteTensor tensor_a, tensor_b; + std::vector shape_a = {1}; + std::vector shape_b = {2}; + InitTensor(shape_a, kTfLiteDynamic, 1.0, &tensor_a); + InitTensor(shape_b, kTfLiteDynamic, 4.0, &tensor_b); + + EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_a)); + EXPECT_TRUE(var.IsInitialized()); + auto* value = var.GetTensor(); + // Variables are always dynamic type. + EXPECT_EQ(kTfLiteDynamic, value->allocation_type); + EXPECT_EQ(kTfLiteFloat32, value->type); + EXPECT_EQ(sizeof(float), value->bytes); + EXPECT_EQ(1, value->dims->size); + EXPECT_EQ(1, value->dims->data[0]); + EXPECT_EQ(1.0f, value->data.f[0]); + + // Second AssignFrom but now tensor_b has different size from the variable. + EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_b)); + EXPECT_TRUE(var.IsInitialized()); + value = var.GetTensor(); + // Variables are always dynamic type. + EXPECT_EQ(kTfLiteDynamic, value->allocation_type); + EXPECT_EQ(kTfLiteFloat32, value->type); + EXPECT_EQ(sizeof(float) * 2, value->bytes); + EXPECT_EQ(1, value->dims->size); + EXPECT_EQ(2, value->dims->data[0]); + EXPECT_EQ(4.0f, value->data.f[0]); + + // Cleanup + TfLiteTensorFree(&tensor_a); + TfLiteTensorFree(&tensor_b); +} + +} // namespace resource +} // namespace tflite