def test_save_load_hdf5_pathlib(self): if sys.version_info < (3, 6): self.skipTest( 'pathlib is only available for python version >= 3.6') path = pathlib.Path(self.get_temp_dir()) / 'model' save.save_model(self.model, path, save_format='h5') save.load_model(path)
def test_save_model_with_dynamic_loss_scaling(self, strategy_fn, h5=False): # TODO(reedwm): Support and test saving model with a mixed_[b]float16 policy # as well. strategy = strategy_fn() if (isinstance(strategy, tf.distribute.MirroredStrategy) and not tf.executing_eagerly()): # TODO(b/121381184): Enable running the test in this case. return # Create and run model. with strategy.scope(): x = layers.Input(shape=(2, ), batch_size=2, dtype=tf.float32) y = mp_test_util.MultiplyLayer()(x) model = models.Model(inputs=x, outputs=y) opt = gradient_descent.SGD(1.) opt = loss_scale_optimizer.LossScaleOptimizer( opt, initial_scale=1., dynamic_growth_steps=2.) model.compile(optimizer=opt, loss='mse', run_eagerly=test_utils.should_run_eagerly()) # Run for 3 steps (6 examples with a batch size of 2) model.fit(np.ones((6, 2)), np.zeros((6, 2)), batch_size=2) self.assertEqual(backend.get_value(opt.loss_scale), 2) self.assertEqual(backend.get_value(opt.dynamic_counter), 1) (weight, ) = model.trainable_weights orig_weight = backend.get_value(weight) # Save model weights. save_path = os.path.join(self.get_temp_dir(), 'model') model.save(save_path, save_format='h5' if h5 else 'tf') # Run model again for 1 step (2 examples with a batch size of 2) model.fit(np.ones((2, 2)), np.zeros((2, 2)), batch_size=2) new_weight = backend.get_value(weight) self.assertNotEqual(new_weight, orig_weight) self.assertEqual(backend.get_value(opt.loss_scale), 4) self.assertEqual(backend.get_value(opt.dynamic_counter), 0) # Load model weights and ensure loss scale weights are restored. model = save.load_model( save_path, custom_objects={'MultiplyLayer': mp_test_util.MultiplyLayer}) (weight, ) = model.trainable_weights loaded_weight = backend.get_value(weight) self.assertEqual(loaded_weight, orig_weight) # Currently the loss scale isn't always saved when the model is saved with # Model.save(). So we assert the loss scale either has the value when it was # saved, or the value it was initialized with. # TODO(reedwm): Always save/restore the loss scale with Model.save(). self.assertIn(backend.get_value(model.optimizer.loss_scale), (1, 2)) self.assertIn(backend.get_value(model.optimizer.dynamic_counter), (0, 1)) # Test optimizer attributes and type self.assertEqual(model.optimizer.initial_scale, 1.) self.assertEqual(model.optimizer.dynamic_growth_steps, 2.) self.assertEqual(type(model.optimizer), loss_scale_optimizer.LossScaleOptimizer)
def _load_and_run_model(self, distribution, saved_dir, predict_dataset, output_name='output_1'): restored_keras_model = save.load_model(saved_dir) return restored_keras_model.predict(predict_dataset, steps=test_base.PREDICT_STEPS)
def _test_saving(self, model, dataset, save_format, use_regularizer): # Save and load model, asserting variable does not change save_path = os.path.join(self.get_temp_dir(), "model") model.save(save_path, save_format=save_format) model = save.load_model(save_path) (layer,) = ( layer for layer in model.layers if "MultiplyLayer" in layer.__class__.__name__ ) expected = 1 - 2**-14 if use_regularizer: expected -= 2 * 2**-14 self.assertEqual(backend.eval(layer.v), expected) # Continue training, and assert variable is correct value model.fit(dataset) new_expected = expected - 2**-14 if use_regularizer: new_expected -= 2 * 2**-14 self.assertEqual(backend.eval(layer.v), new_expected) # Load saved model again, and assert variable is previous value model = save.load_model(save_path) (layer,) = ( layer for layer in model.layers if "MultiplyLayer" in layer.__class__.__name__ ) self.assertEqual(backend.eval(layer.v), expected) # Ensure various dtype-related aspects of the layer are correct self.assertEqual(layer.dtype, "float32") self.assertEqual(layer.dtype_policy.name, "mixed_float16") self.assertEqual(layer.v.dtype, "float32") self.assertEqual(layer(np.ones((2, 1))).dtype, "float16") self.assertEqual(type(model.dtype_policy), policy.Policy) self.assertEqual( layer.get_config()["dtype"], {"class_name": "Policy", "config": {"name": "mixed_float16"}}, )
def _test_saving(self, model, dataset, save_format, use_regularizer): # Save and load model, asserting variable does not change save_path = os.path.join(self.get_temp_dir(), 'model') model.save(save_path, save_format=save_format) model = save.load_model(save_path) (layer, ) = (layer for layer in model.layers if 'MultiplyLayer' in layer.__class__.__name__) expected = 1 - 2**-14 if use_regularizer: expected -= 2 * 2**-14 self.assertEqual(backend.eval(layer.v), expected) # Continue training, and assert variable is correct value model.fit(dataset) new_expected = expected - 2**-14 if use_regularizer: new_expected -= 2 * 2**-14 self.assertEqual(backend.eval(layer.v), new_expected) # Load saved model again, and assert variable is previous value model = save.load_model(save_path) (layer, ) = (layer for layer in model.layers if 'MultiplyLayer' in layer.__class__.__name__) self.assertEqual(backend.eval(layer.v), expected) # Ensure various dtype-related aspects of the layer are correct self.assertEqual(layer.dtype, 'float32') self.assertEqual( get_layer_policy.get_layer_policy(layer).name, 'mixed_float16') self.assertEqual(layer.v.dtype, 'float32') self.assertEqual(layer(np.ones((2, 1))).dtype, 'float16') # Loading a model always loads with a v2 Policy, even if saved with a # PolicyV1. self.assertEqual(type(model.dtype_policy), policy.Policy) self.assertEqual(layer.get_config()['dtype'], { 'class_name': 'Policy', 'config': { 'name': 'mixed_float16' } })
def test_restore_old_saved_model(self): saved_model_dir = os.path.join( flags.FLAGS['test_srcdir'].value, 'org_keras/keras', 'mixed_precision/testdata/lso_savedmodel_tf2.2') # saved_model_dir = test.test_src_dir_path( # 'python/keras/mixed_precision/testdata/' # 'lso_savedmodel_tf2.2') model = save.load_model(saved_model_dir) expected_kernel = np.array([[9.229685, 10.901115], [10.370763, 9.757362]]) self.assertAllClose(backend.eval(model.weights[0]), expected_kernel) self.assertEqual(type(model.optimizer), loss_scale_optimizer.LossScaleOptimizer)
def test_state_saving_and_loading(self): with self.cached_session(): input_data = np.random.random((1, 2)) rff_layer = kernel_layers.RandomFourierFeatures(output_dim=10, scale=3.0) inputs = input_layer.Input((2,)) outputs = rff_layer(inputs) model = training.Model(inputs, outputs) output_data = model.predict(input_data) temp_dir = self.get_temp_dir() self.addCleanup(shutil.rmtree, temp_dir) saved_model_dir = os.path.join(temp_dir, 'rff_model') model.save(saved_model_dir) new_model = save.load_model(saved_model_dir) new_output_data = new_model.predict(input_data) self.assertAllClose(output_data, new_output_data, atol=1e-4)
def test_saving_model_with_custom_object(self): with generic_utils.custom_object_scope(), self.cached_session(): @generic_utils.register_keras_serializable() class CustomLoss(losses.MeanSquaredError): pass model = sequential.Sequential( [core.Dense(units=1, input_shape=(1,))]) model.compile(optimizer='sgd', loss=CustomLoss()) model.fit(np.zeros([10, 1]), np.zeros([10, 1])) temp_dir = self.get_temp_dir() filepath = os.path.join(temp_dir, 'saving') model.save(filepath) # Make sure the model can be correctly load back. _ = save.load_model(filepath, compile=True)
def deserialize_model_from_bytecode(serialized_model): """Reconstruct a Model from the output of `serialize_model_as_bytecode`. Args: serialized_model: (np.array) return value from `serialize_model_as_bytecode`. Returns: keras.Model: Keras Model instance. """ temp_dir = f"ram://{uuid.uuid4()}" b = io.BytesIO(serialized_model) with tarfile.open(fileobj=b, mode="r") as archive: for name in archive.getnames(): dest_path = tf.io.gfile.join(temp_dir, name) member = archive.getmember(name) tf.io.gfile.makedirs(os.path.dirname(dest_path)) if member.isfile(): with tf.io.gfile.GFile(dest_path, "wb") as f: f.write(archive.extractfile(name).read()) model = save_module.load_model(temp_dir) tf.io.gfile.rmtree(temp_dir) return model
def test_save_load_hdf5_pathlib(self): path = pathlib.Path(self.get_temp_dir()) / 'model' save.save_model(self.model, path, save_format='h5') save.load_model(path)
def test_load_file_not_found(self): path = pathlib.Path(self.get_temp_dir()) / 'does_not_exist' with self.assertRaisesRegex(IOError, 'No file or directory found at'): save.load_model(path)
def test_save_load_tf_string(self): path = os.path.join(self.get_temp_dir(), 'model') save.save_model(self.model, path, save_format='tf') save.load_model(path)