1# Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15# pylint: disable=protected-access 16"""Tests for saving/loading function for keras Model.""" 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import os 22import shutil 23 24from absl.testing import parameterized 25import numpy as np 26 27from tensorflow.python import keras 28from tensorflow.python import tf2 29from tensorflow.python.client import session 30from tensorflow.python.eager import context 31from tensorflow.python.framework import dtypes 32from tensorflow.python.framework import ops 33from tensorflow.python.framework import tensor_spec 34from tensorflow.python.framework import test_util 35from tensorflow.python.keras.engine import training 36from tensorflow.python.keras.optimizer_v2 import adadelta 37from tensorflow.python.keras.saving import saved_model as keras_saved_model 38from tensorflow.python.keras.utils import mode_keys 39from tensorflow.python.keras.utils import tf_utils 40from tensorflow.python.ops import array_ops 41from tensorflow.python.platform import test 42from tensorflow.python.saved_model import loader_impl 43from tensorflow.python.saved_model import model_utils 44from tensorflow.python.training import training as training_module 45 46 47class TestModelSavingandLoading(test.TestCase): 48 49 def _save_model_dir(self, dirname='saved_model'): 50 temp_dir = self.get_temp_dir() 51 self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True) 52 return os.path.join(temp_dir, dirname) 53 54 def test_saving_sequential_model(self): 55 with self.cached_session(): 56 model = keras.models.Sequential() 57 model.add(keras.layers.Dense(2, input_shape=(3,))) 58 model.add(keras.layers.RepeatVector(3)) 59 model.add(keras.layers.TimeDistributed(keras.layers.Dense(3))) 60 model.compile( 61 loss=keras.losses.MSE, 62 optimizer=keras.optimizers.RMSprop(lr=0.0001), 63 metrics=[keras.metrics.categorical_accuracy], 64 sample_weight_mode='temporal') 65 x = np.random.random((1, 3)) 66 y = np.random.random((1, 3, 3)) 67 model.train_on_batch(x, y) 68 69 ref_y = model.predict(x) 70 71 saved_model_dir = self._save_model_dir() 72 keras_saved_model.export_saved_model(model, saved_model_dir) 73 74 loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir) 75 y = loaded_model.predict(x) 76 self.assertAllClose(ref_y, y, atol=1e-05) 77 78 @test_util.run_in_graph_and_eager_modes 79 def test_saving_sequential_model_without_compile(self): 80 with self.cached_session(): 81 model = keras.models.Sequential() 82 model.add(keras.layers.Dense(2, input_shape=(3,))) 83 model.add(keras.layers.RepeatVector(3)) 84 model.add(keras.layers.TimeDistributed(keras.layers.Dense(3))) 85 86 x = np.random.random((1, 3)) 87 ref_y = model.predict(x) 88 89 saved_model_dir = self._save_model_dir() 90 keras_saved_model.export_saved_model(model, saved_model_dir) 91 loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir) 92 93 y = loaded_model.predict(x) 94 self.assertAllClose(ref_y, y, atol=1e-05) 95 96 def test_saving_functional_model(self): 97 with self.cached_session(): 98 inputs = keras.layers.Input(shape=(3,)) 99 x = keras.layers.Dense(2)(inputs) 100 output = keras.layers.Dense(3)(x) 101 102 model = keras.models.Model(inputs, output) 103 model.compile( 104 loss=keras.losses.MSE, 105 optimizer=keras.optimizers.RMSprop(lr=0.0001), 106 metrics=[keras.metrics.categorical_accuracy]) 107 x = np.random.random((1, 3)) 108 y = np.random.random((1, 3)) 109 model.train_on_batch(x, y) 110 111 ref_y = model.predict(x) 112 113 saved_model_dir = self._save_model_dir() 114 keras_saved_model.export_saved_model(model, saved_model_dir) 115 loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir) 116 117 y = loaded_model.predict(x) 118 self.assertAllClose(ref_y, y, atol=1e-05) 119 120 @test_util.run_in_graph_and_eager_modes 121 def test_saving_functional_model_without_compile(self): 122 with self.cached_session(): 123 inputs = keras.layers.Input(shape=(3,)) 124 x = keras.layers.Dense(2)(inputs) 125 output = keras.layers.Dense(3)(x) 126 127 model = keras.models.Model(inputs, output) 128 129 x = np.random.random((1, 3)) 130 y = np.random.random((1, 3)) 131 132 ref_y = model.predict(x) 133 134 saved_model_dir = self._save_model_dir() 135 keras_saved_model.export_saved_model(model, saved_model_dir) 136 loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir) 137 138 y = loaded_model.predict(x) 139 self.assertAllClose(ref_y, y, atol=1e-05) 140 141 @test_util.run_in_graph_and_eager_modes 142 def test_saving_with_tf_optimizer(self): 143 model = keras.models.Sequential() 144 model.add(keras.layers.Dense(2, input_shape=(3,))) 145 model.add(keras.layers.Dense(3)) 146 model.compile( 147 loss='mse', 148 optimizer=training_module.RMSPropOptimizer(0.1), 149 metrics=['acc']) 150 151 x = np.random.random((1, 3)) 152 y = np.random.random((1, 3)) 153 model.train_on_batch(x, y) 154 ref_y = model.predict(x) 155 156 saved_model_dir = self._save_model_dir() 157 keras_saved_model.export_saved_model(model, saved_model_dir) 158 loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir) 159 loaded_model.compile( 160 loss='mse', 161 optimizer=training_module.RMSPropOptimizer(0.1), 162 metrics=['acc']) 163 y = loaded_model.predict(x) 164 self.assertAllClose(ref_y, y, atol=1e-05) 165 166 # test that new updates are the same with both models 167 x = np.random.random((1, 3)) 168 y = np.random.random((1, 3)) 169 170 ref_loss = model.train_on_batch(x, y) 171 loss = loaded_model.train_on_batch(x, y) 172 self.assertAllClose(ref_loss, loss, atol=1e-05) 173 174 ref_y = model.predict(x) 175 y = loaded_model.predict(x) 176 self.assertAllClose(ref_y, y, atol=1e-05) 177 178 # test saving/loading again 179 saved_model_dir2 = self._save_model_dir('saved_model_2') 180 keras_saved_model.export_saved_model(loaded_model, saved_model_dir2) 181 loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir2) 182 y = loaded_model.predict(x) 183 self.assertAllClose(ref_y, y, atol=1e-05) 184 185 def test_saving_subclassed_model_raise_error(self): 186 # For now, saving subclassed model should raise an error. It should be 187 # avoided later with loading from SavedModel.pb. 188 189 class SubclassedModel(training.Model): 190 191 def __init__(self): 192 super(SubclassedModel, self).__init__() 193 self.layer1 = keras.layers.Dense(3) 194 self.layer2 = keras.layers.Dense(1) 195 196 def call(self, inp): 197 return self.layer2(self.layer1(inp)) 198 199 model = SubclassedModel() 200 201 saved_model_dir = self._save_model_dir() 202 with self.assertRaises(NotImplementedError): 203 keras_saved_model.export_saved_model(model, saved_model_dir) 204 205 206class LayerWithLearningPhase(keras.engine.base_layer.Layer): 207 208 def call(self, x): 209 phase = keras.backend.learning_phase() 210 output = tf_utils.smart_cond( 211 phase, lambda: x * 0, lambda: array_ops.identity(x)) 212 if not context.executing_eagerly(): 213 output._uses_learning_phase = True # pylint: disable=protected-access 214 return output 215 216 def compute_output_shape(self, input_shape): 217 return input_shape 218 219 220def functional_model(uses_learning_phase=True): 221 inputs = keras.layers.Input(shape=(3,)) 222 x = keras.layers.Dense(2)(inputs) 223 x = keras.layers.Dense(3)(x) 224 if uses_learning_phase: 225 x = LayerWithLearningPhase()(x) 226 return keras.models.Model(inputs, x) 227 228 229def sequential_model(uses_learning_phase=True): 230 model = keras.models.Sequential() 231 model.add(keras.layers.Dense(2, input_shape=(3,))) 232 model.add(keras.layers.Dense(3)) 233 if uses_learning_phase: 234 model.add(LayerWithLearningPhase()) 235 return model 236 237 238def sequential_model_without_input_shape(uses_learning_phase=True): 239 model = keras.models.Sequential() 240 model.add(keras.layers.Dense(2)) 241 model.add(keras.layers.Dense(3)) 242 if uses_learning_phase: 243 model.add(LayerWithLearningPhase()) 244 return model 245 246 247class Subclassed(keras.models.Model): 248 249 def __init__(self): 250 super(Subclassed, self).__init__() 251 self.dense1 = keras.layers.Dense(2) 252 self.dense2 = keras.layers.Dense(3) 253 254 def call(self, inputs): 255 x = self.dense1(inputs) 256 x = self.dense2(x) 257 return x 258 259 260def subclassed_model(): 261 return Subclassed() 262 263 264def load_model(sess, path, mode): 265 tags = model_utils.EXPORT_TAG_MAP[mode] 266 sig_def_key = model_utils.SIGNATURE_KEY_MAP[mode] 267 268 meta_graph_def = loader_impl.load(sess, tags, path) 269 inputs = { 270 k: sess.graph.get_tensor_by_name(v.name) 271 for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()} 272 outputs = { 273 k: sess.graph.get_tensor_by_name(v.name) 274 for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()} 275 return inputs, outputs, meta_graph_def 276 277 278@test_util.run_all_in_graph_and_eager_modes 279class TestModelSavedModelExport(test.TestCase, parameterized.TestCase): 280 281 def _save_model_dir(self, dirname='saved_model'): 282 temp_dir = self.get_temp_dir() 283 self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True) 284 return os.path.join(temp_dir, dirname) 285 286 @parameterized.parameters( 287 { 288 'model_builder': functional_model, 289 'uses_learning_phase': True, 290 'optimizer_cls': adadelta.Adadelta, 291 'train_before_export': True}, 292 { 293 'model_builder': functional_model, 294 'uses_learning_phase': True, 295 'optimizer_cls': training_module.AdadeltaOptimizer, 296 'train_before_export': False}, 297 { 298 'model_builder': functional_model, 299 'uses_learning_phase': False, 300 'optimizer_cls': None, 301 'train_before_export': False}, 302 { 303 'model_builder': sequential_model, 304 'uses_learning_phase': True, 305 'optimizer_cls': training_module.AdadeltaOptimizer, 306 'train_before_export': True}, 307 { 308 'model_builder': sequential_model, 309 'uses_learning_phase': True, 310 'optimizer_cls': adadelta.Adadelta, 311 'train_before_export': False}, 312 { 313 'model_builder': sequential_model, 314 'uses_learning_phase': False, 315 'optimizer_cls': None, 316 'train_before_export': False}, 317 { 318 'model_builder': sequential_model_without_input_shape, 319 'uses_learning_phase': True, 320 'optimizer_cls': training_module.AdadeltaOptimizer, 321 'train_before_export': False}) 322 def testSaveAndLoadSavedModelExport( 323 self, model_builder, uses_learning_phase, optimizer_cls, 324 train_before_export): 325 optimizer = None if optimizer_cls is None else optimizer_cls() 326 327 saved_model_dir = self._save_model_dir() 328 329 np.random.seed(130) 330 input_arr = np.random.random((1, 3)) 331 target_arr = np.random.random((1, 3)) 332 333 model = model_builder(uses_learning_phase) 334 if optimizer is not None: 335 model.compile( 336 loss='mse', 337 optimizer=optimizer, 338 metrics=['mae']) 339 if train_before_export: 340 model.train_on_batch(input_arr, target_arr) 341 342 ref_loss, ref_mae = model.evaluate(input_arr, target_arr) 343 344 ref_predict = model.predict(input_arr) 345 346 # Export SavedModel 347 keras_saved_model.export_saved_model(model, saved_model_dir) 348 349 input_name = model.input_names[0] 350 output_name = model.output_names[0] 351 target_name = output_name + '_target' 352 353 # Load predict graph, and test predictions 354 with session.Session(graph=ops.Graph()) as sess: 355 inputs, outputs, _ = load_model(sess, saved_model_dir, 356 mode_keys.ModeKeys.PREDICT) 357 358 predictions = sess.run(outputs[output_name], 359 {inputs[input_name]: input_arr}) 360 self.assertAllClose(ref_predict, predictions, atol=1e-05) 361 362 if optimizer: 363 # Load eval graph, and test predictions, loss and metric values 364 with session.Session(graph=ops.Graph()) as sess: 365 inputs, outputs, _ = load_model(sess, saved_model_dir, 366 mode_keys.ModeKeys.TEST) 367 368 # First obtain the loss and predictions, and run the metric update op by 369 # feeding in the inputs and targets. 370 metrics_name = 'mae' if tf2.enabled() else 'mean_absolute_error' 371 metrics_update_op_key = 'metrics/' + metrics_name + '/update_op' 372 metrics_value_op_key = 'metrics/' + metrics_name + '/value' 373 374 loss, predictions, _ = sess.run( 375 (outputs['loss'], outputs['predictions/' + output_name], 376 outputs[metrics_update_op_key]), { 377 inputs[input_name]: input_arr, 378 inputs[target_name]: target_arr 379 }) 380 381 # The metric value should be run after the update op, to ensure that it 382 # reflects the correct value. 383 metric_value = sess.run(outputs[metrics_value_op_key]) 384 385 self.assertEqual(int(train_before_export), 386 sess.run(training_module.get_global_step())) 387 self.assertAllClose(ref_loss, loss, atol=1e-05) 388 self.assertAllClose(ref_mae, metric_value, atol=1e-05) 389 self.assertAllClose(ref_predict, predictions, atol=1e-05) 390 391 # Load train graph, and check for the train op, and prediction values 392 with session.Session(graph=ops.Graph()) as sess: 393 inputs, outputs, meta_graph_def = load_model( 394 sess, saved_model_dir, mode_keys.ModeKeys.TRAIN) 395 self.assertEqual(int(train_before_export), 396 sess.run(training_module.get_global_step())) 397 self.assertIn('loss', outputs) 398 self.assertIn(metrics_update_op_key, outputs) 399 self.assertIn(metrics_value_op_key, outputs) 400 self.assertIn('predictions/' + output_name, outputs) 401 402 # Train for a step 403 train_op = loader_impl.get_train_op(meta_graph_def) 404 train_outputs, _ = sess.run( 405 [outputs, train_op], {inputs[input_name]: input_arr, 406 inputs[target_name]: target_arr}) 407 self.assertEqual(int(train_before_export) + 1, 408 sess.run(training_module.get_global_step())) 409 410 if uses_learning_phase: 411 self.assertAllClose( 412 [[0, 0, 0]], train_outputs['predictions/' + output_name], 413 atol=1e-05) 414 else: 415 self.assertNotAllClose( 416 [[0, 0, 0]], train_outputs['predictions/' + output_name], 417 atol=1e-05) 418 419 def testSaveAndLoadSavedModelWithCustomObject(self): 420 saved_model_dir = self._save_model_dir() 421 with session.Session(graph=ops.Graph()) as sess: 422 def relu6(x): 423 return keras.backend.relu(x, max_value=6) 424 inputs = keras.layers.Input(shape=(1,)) 425 outputs = keras.layers.Activation(relu6)(inputs) 426 model = keras.models.Model(inputs, outputs) 427 keras_saved_model.export_saved_model( 428 model, saved_model_dir, custom_objects={'relu6': relu6}) 429 with session.Session(graph=ops.Graph()) as sess: 430 inputs, outputs, _ = load_model(sess, saved_model_dir, 431 mode_keys.ModeKeys.PREDICT) 432 input_name = model.input_names[0] 433 output_name = model.output_names[0] 434 predictions = sess.run( 435 outputs[output_name], {inputs[input_name]: [[7], [-3], [4]]}) 436 self.assertAllEqual([[6], [0], [4]], predictions) 437 438 def testAssertModelCloneSameObjectsIgnoreOptimizer(self): 439 input_arr = np.random.random((1, 3)) 440 target_arr = np.random.random((1, 3)) 441 442 model_graph = ops.Graph() 443 clone_graph = ops.Graph() 444 445 # Create two models with the same layers but different optimizers. 446 with session.Session(graph=model_graph): 447 inputs = keras.layers.Input(shape=(3,)) 448 x = keras.layers.Dense(2)(inputs) 449 x = keras.layers.Dense(3)(x) 450 model = keras.models.Model(inputs, x) 451 452 model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer()) 453 model.train_on_batch(input_arr, target_arr) 454 455 with session.Session(graph=clone_graph): 456 inputs = keras.layers.Input(shape=(3,)) 457 x = keras.layers.Dense(2)(inputs) 458 x = keras.layers.Dense(3)(x) 459 clone = keras.models.Model(inputs, x) 460 clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001)) 461 clone.train_on_batch(input_arr, target_arr) 462 463 keras_saved_model._assert_same_non_optimizer_objects( 464 model, model_graph, clone, clone_graph) 465 466 def testAssertModelCloneSameObjectsThrowError(self): 467 input_arr = np.random.random((1, 3)) 468 target_arr = np.random.random((1, 3)) 469 470 model_graph = ops.Graph() 471 clone_graph = ops.Graph() 472 473 # Create two models with the same layers but different optimizers. 474 with session.Session(graph=model_graph): 475 inputs = keras.layers.Input(shape=(3,)) 476 x = keras.layers.Dense(2)(inputs) 477 x = keras.layers.Dense(3)(x) 478 model = keras.models.Model(inputs, x) 479 480 model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer()) 481 model.train_on_batch(input_arr, target_arr) 482 483 with session.Session(graph=clone_graph): 484 inputs = keras.layers.Input(shape=(3,)) 485 x = keras.layers.Dense(2)(inputs) 486 x = keras.layers.Dense(4)(x) 487 x = keras.layers.Dense(3)(x) 488 clone = keras.models.Model(inputs, x) 489 clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001)) 490 clone.train_on_batch(input_arr, target_arr) 491 492 def testSaveSequentialModelWithoutInputShapes(self): 493 model = sequential_model_without_input_shape(True) 494 # A Sequential model that hasn't been built should raise an error. 495 with self.assertRaisesRegexp(ValueError, 'Please build the model'): 496 keras_saved_model.export_saved_model(model, '') 497 498 saved_model_dir = self._save_model_dir() 499 keras_saved_model.export_saved_model( 500 model, 501 saved_model_dir, 502 input_signature=tensor_spec.TensorSpec( 503 shape=(10, 11, 12, 13, 14), dtype=dtypes.float32, 504 name='spec_input')) 505 506 with session.Session(graph=ops.Graph()) as sess: 507 inputs, outputs, _ = load_model(sess, saved_model_dir, 508 mode_keys.ModeKeys.PREDICT) 509 self.assertEqual(5, inputs[next(iter(inputs.keys()))].shape.ndims) 510 self.assertEqual(5, outputs[next(iter(outputs.keys()))].shape.ndims) 511 self.assertEqual(3, outputs[next(iter(outputs.keys()))].shape[-1]) 512 513 @parameterized.parameters( 514 { 515 'model_builder': sequential_model_without_input_shape, 516 'input_signature': [tensor_spec.TensorSpec(shape=[None, 3], 517 dtype=dtypes.float32)]}, 518 { 519 'model_builder': subclassed_model, 520 'input_signature': [tensor_spec.TensorSpec(shape=[None, 3], 521 dtype=dtypes.float32)]}) 522 def testServingOnly(self, model_builder, input_signature): 523 if context.executing_eagerly(): 524 saved_model_dir = self._save_model_dir() 525 input_arr = np.random.random((5, 3)).astype(np.float32) 526 model = model_builder() 527 ref_predict = model.predict(input_arr) 528 529 keras_saved_model.export_saved_model( 530 model, 531 saved_model_dir, 532 serving_only=True, 533 input_signature=input_signature) 534 535 # Load predict graph, and test predictions 536 with session.Session(graph=ops.Graph()) as sess: 537 inputs, outputs, _ = load_model(sess, saved_model_dir, 538 mode_keys.ModeKeys.PREDICT) 539 predictions = sess.run(outputs[next(iter(outputs.keys()))], 540 {inputs[next(iter(inputs.keys()))]: input_arr}) 541 self.assertAllClose(ref_predict, predictions, atol=1e-05) 542 543 544if __name__ == '__main__': 545 test.main() 546