From b737d47b533ec0abfc12977022f8fb8e550e7f44 Mon Sep 17 00:00:00 2001 From: Scott Zhu Date: Sat, 11 Apr 2020 14:03:45 -0700 Subject: [PATCH] Move keras related convert_to_constants_test to keras/tests. PiperOrigin-RevId: 306056993 Change-Id: I220960a74da3563de424bf50fbac6739b3c5df6a --- .../framework/convert_to_constants_test.py | 77 -------- tensorflow/python/keras/tests/BUILD | 22 +++ .../keras/tests/convert_to_constants_test.py | 181 ++++++++++++++++++ 3 files changed, 203 insertions(+), 77 deletions(-) create mode 100644 tensorflow/python/keras/tests/convert_to_constants_test.py diff --git a/tensorflow/python/framework/convert_to_constants_test.py b/tensorflow/python/framework/convert_to_constants_test.py index e44847e2c6e..56a2e2b13a4 100644 --- a/tensorflow/python/framework/convert_to_constants_test.py +++ b/tensorflow/python/framework/convert_to_constants_test.py @@ -22,7 +22,6 @@ import os import numpy as np -from tensorflow.python import keras from tensorflow.python.client import session as session_lib from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op @@ -34,7 +33,6 @@ from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import cond_v2 from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import rnn from tensorflow.python.ops import rnn_cell_impl @@ -221,29 +219,6 @@ class VariablesToConstantsTest(test.TestCase): input_func) self._testConvertedFunction(root, root.add, output_func, input_data) - @test_util.run_v2_only - def testKerasModel(self): - """Test a basic Keras model with Variables.""" - input_data = {"x": constant_op.constant(1., shape=[1, 1])} - - # Create a simple Keras model. - x = [-1, 0, 1, 2, 3, 4] - y = [-3, -1, 1, 3, 5, 7] - - model = keras.models.Sequential( - [keras.layers.Dense(units=1, input_shape=[1])]) - model.compile(optimizer="sgd", loss="mean_squared_error") - model.fit(x, y, epochs=1) - - @def_function.function(input_signature=[ - tensor_spec.TensorSpec(shape=[1, 1], dtype=dtypes.float32) - ]) - def to_save(x): - return model(x) - - root, output_func = self._freezeModel(to_save) - self._testConvertedFunction(root, root.f, output_func, input_data) - def _singleMetaGraphSavedModel(self): export_graph = ops.Graph() with export_graph.as_default(): @@ -413,58 +388,6 @@ class VariablesToConstantsTest(test.TestCase): root, output_func = self._freezeModel(model) self._testConvertedFunction(root, root.f, output_func, input_data) - @test_util.run_v2_only - def testKerasLSTM(self): - """Test a Keras LSTM containing dynamic_rnn ops.""" - input_data = { - "x": - constant_op.constant( - np.array( - np.random.random_sample((10, 10, 10)), dtype=np.float32)) - } - - model = keras.models.Sequential( - [keras.layers.LSTM(units=10, input_shape=(10, 10))]) - - @def_function.function(input_signature=[ - tensor_spec.TensorSpec(shape=[10, 10, 10], dtype=dtypes.float32) - ]) - def to_save(x): - return model(x) - - root, output_func = self._freezeModel(to_save) - self._testConvertedFunction(root, root.f, output_func, input_data) - - @test_util.run_v2_only - def testEmbeddings(self): - """Test model with embeddings.""" - input_data = { - "x": - constant_op.constant( - np.array(np.random.random_sample((20)), dtype=np.int32)) - } - - class EmbeddingModel(keras.Model): - - def __init__(self): - super(EmbeddingModel, self).__init__() - self.shared_weights = self.add_weight( - "weights", - shape=(2000, 300), - dtype=dtypes.float32, - initializer=init_ops.random_normal_initializer( - mean=0.0, stddev=300**(-0.5))) - - @def_function.function(input_signature=[ - tensor_spec.TensorSpec(shape=(20), dtype=dtypes.int32) - ]) - def func(self, x): - return array_ops.gather(self.shared_weights, x) - - model = EmbeddingModel() - root, output_func = self._freezeModel(model.func) - self._testConvertedFunction(root, root.f, output_func, input_data) - if __name__ == "__main__": test.main() diff --git a/tensorflow/python/keras/tests/BUILD b/tensorflow/python/keras/tests/BUILD index b3e6503c58c..0caa486bc42 100644 --- a/tensorflow/python/keras/tests/BUILD +++ b/tensorflow/python/keras/tests/BUILD @@ -26,6 +26,28 @@ tf_py_test( ], ) +tf_py_test( + name = "convert_to_constants_test", + srcs = ["convert_to_constants_test.py"], + python_version = "PY3", + shard_count = 4, + deps = [ + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:convert_to_constants", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:init_ops", + "//tensorflow/python:tensor_spec", + "//tensorflow/python:util", + "//tensorflow/python/eager:def_function", + "//tensorflow/python/saved_model:load", + "//tensorflow/python/saved_model:save", + "//tensorflow/python/training/tracking", + ], +) + tf_py_test( name = "custom_training_loop_test", srcs = ["custom_training_loop_test.py"], diff --git a/tensorflow/python/keras/tests/convert_to_constants_test.py b/tensorflow/python/keras/tests/convert_to_constants_test.py new file mode 100644 index 00000000000..21081682089 --- /dev/null +++ b/tensorflow/python/keras/tests/convert_to_constants_test.py @@ -0,0 +1,181 @@ +# Copyright 2019 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. +# ============================================================================== +"""Tests for convert_to_constants.py.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +import numpy as np + +from tensorflow.python import keras +from tensorflow.python.eager import def_function +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import convert_to_constants +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_spec +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import init_ops +from tensorflow.python.platform import test +from tensorflow.python.saved_model.load import load +from tensorflow.python.saved_model.save import save +from tensorflow.python.training.tracking import tracking +from tensorflow.python.util import nest + + +class VariablesToConstantsTest(test.TestCase): + + def _freezeModel(self, model): + """Freezes the model. + + Args: + model: Function. + + Returns: + root: AutoTrackable object with original ConcreteFunction. + output_func: frozen ConcreteFunction. + """ + root = tracking.AutoTrackable() + root.f = model + input_func = root.f.get_concrete_function() + + output_func = convert_to_constants.convert_variables_to_constants_v2( + input_func, lower_control_flow=False) + return root, output_func + + def _hasStatefulPartitionedCallOp(self, graph_def): + """Determines if a StatefulPartitionedCall op exists in the graph.""" + for node in graph_def.node: + if node.op == "StatefulPartitionedCall": + return True + return False + + def _getNumVariables(self, graph_def): + """Returns the number of ReadVariableOp in the graph.""" + return sum(node.op == "ReadVariableOp" for node in graph_def.node) + + def _testConvertedFunction(self, obj, func, converted_concrete_func, + input_data): + # Ensure the converted graph has no variables and no function calls. + constant_graph_def = converted_concrete_func.graph.as_graph_def() + self.assertEqual(0, self._getNumVariables(constant_graph_def)) + self.assertFalse(self._hasStatefulPartitionedCallOp(constant_graph_def)) + + # Check that the converted ConcreteFunction produces the same result as the + # original Function. + expected_value = nest.flatten(func(**input_data)) + actual_value = nest.flatten(converted_concrete_func(**input_data)) + + for expected, actual in zip(expected_value, actual_value): + np.testing.assert_almost_equal(expected.numpy(), actual.numpy()) + + # Ensure the shape is retained. + for tensor in converted_concrete_func.inputs: + actual_shape = input_data[tensor.name.split(":")[0]].shape + self.assertEqual(tensor.shape, actual_shape) + + # Save the converted ConcreteFunction as a signature. + save_dir = os.path.join(self.get_temp_dir(), "frozen_saved_model") + root = tracking.AutoTrackable() + root.f = converted_concrete_func + save(root, save_dir, {"mykey": converted_concrete_func}) + + # Load it back and make sure it works. + loaded_obj = load(save_dir) + actual_value = nest.flatten(loaded_obj.signatures["mykey"](**input_data)) + for expected, actual in zip(expected_value, actual_value): + np.testing.assert_almost_equal(expected.numpy(), actual.numpy()) + + @test_util.run_v2_only + def testKerasModel(self): + """Test a basic Keras model with Variables.""" + input_data = {"x": constant_op.constant(1., shape=[1, 1])} + + # Create a simple Keras model. + x = [-1, 0, 1, 2, 3, 4] + y = [-3, -1, 1, 3, 5, 7] + + model = keras.models.Sequential( + [keras.layers.Dense(units=1, input_shape=[1])]) + model.compile(optimizer="sgd", loss="mean_squared_error") + model.fit(x, y, epochs=1) + + @def_function.function(input_signature=[ + tensor_spec.TensorSpec(shape=[1, 1], dtype=dtypes.float32) + ]) + def to_save(x): + return model(x) + + root, output_func = self._freezeModel(to_save) + self._testConvertedFunction(root, root.f, output_func, input_data) + + @test_util.run_v2_only + def testKerasLSTM(self): + """Test a Keras LSTM containing dynamic_rnn ops.""" + input_data = { + "x": + constant_op.constant( + np.array( + np.random.random_sample((10, 10, 10)), dtype=np.float32)) + } + + model = keras.models.Sequential( + [keras.layers.LSTM(units=10, input_shape=(10, 10))]) + + @def_function.function(input_signature=[ + tensor_spec.TensorSpec(shape=[10, 10, 10], dtype=dtypes.float32) + ]) + def to_save(x): + return model(x) + + root, output_func = self._freezeModel(to_save) + self._testConvertedFunction(root, root.f, output_func, input_data) + + @test_util.run_v2_only + def testEmbeddings(self): + """Test model with embeddings.""" + input_data = { + "x": + constant_op.constant( + np.array(np.random.random_sample((20)), dtype=np.int32)) + } + + class EmbeddingModel(keras.Model): + + def __init__(self): + super(EmbeddingModel, self).__init__() + self.shared_weights = self.add_weight( + "weights", + shape=(2000, 300), + dtype=dtypes.float32, + initializer=init_ops.random_normal_initializer( + mean=0.0, stddev=300**(-0.5))) + + @def_function.function(input_signature=[ + tensor_spec.TensorSpec(shape=(20), dtype=dtypes.int32) + ]) + def func(self, x): + return array_ops.gather(self.shared_weights, x) + + model = EmbeddingModel() + root, output_func = self._freezeModel(model.func) + self._testConvertedFunction(root, root.f, output_func, input_data) + + +if __name__ == "__main__": + test.main()