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