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