def test_serializing_model_with_loss_with_custom_object_scope(self, value): with generic_utils.custom_object_scope({ 'MyMeanAbsoluteError': MyMeanAbsoluteError, 'my_mae': my_mae, 'Bias': testing_utils.Bias, }): model = _get_multi_io_model() model.compile(optimizer_v2.gradient_descent.SGD(0.1), loss=value, run_eagerly=testing_utils.should_run_eagerly()) history = model.fit([self.x, self.x], [self.y, self.y], batch_size=3, epochs=3, sample_weight=[self.w, self.w]) # Assert training. self.assertAllClose(history.history['loss'], [2., 1.6, 1.2], 1e-3) eval_results = model.evaluate([self.x, self.x], [self.y, self.y], sample_weight=[self.w, self.w]) if h5py is None: return model.save(self.model_filename) loaded_model = keras.models.load_model(self.model_filename) loaded_model.predict([self.x, self.x]) loaded_eval_results = loaded_model.evaluate( [self.x, self.x], [self.y, self.y], sample_weight=[self.w, self.w]) # Assert all evaluation results are the same. self.assertAllClose(eval_results, loaded_eval_results, 1e-9)
def quantize_scope(*args): """Provides a scope in which Quantized layers and models can be deserialized. If a keras model or layer has been quantized, it needs to be within this scope to be successfully deserialized. Args: *args: Variable length list of dictionaries of name, class pairs to add to the scope created by this method. Returns: Object of type `CustomObjectScope` with quantization objects included. Example: ```python keras.models.save_model(quantized_model, keras_file) with quantize_scope(): loaded_model = keras.models.load_model(keras_file) ``` """ quantization_objects = { 'QuantizeAnnotate': quantize_annotate_mod.QuantizeAnnotate, 'QuantizeAwareActivation': quantize_aware_activation.QuantizeAwareActivation, 'QuantizeWrapper': quantize_wrapper.QuantizeWrapper, # TODO(tf-mot): add way for different quantization schemes to modify this. '_DepthwiseConvBatchNorm2D': conv_batchnorm._DepthwiseConvBatchNorm2D, # pylint: disable=protected-access '_ConvBatchNorm2D': conv_batchnorm._ConvBatchNorm2D # pylint: disable=protected-access } quantization_objects.update(tflite_quantize_registry._types_dict()) # pylint: disable=protected-access quantization_objects.update(quantizers._types_dict()) # pylint: disable=protected-access return custom_object_scope(*(args + (quantization_objects, )))
def test_serialize_custom_model_compile(self): with generic_utils.custom_object_scope(): @generic_utils.register_keras_serializable(package="dummy-package") class DummySparseCategoricalCrossentropyLoss(losses.LossFunctionWrapper): # This loss is identical equal to tf.keras.losses.SparseCategoricalCrossentropy def __init__( self, from_logits=False, reduction=losses_utils.ReductionV2.AUTO, name="dummy_sparse_categorical_crossentropy_loss", ): super(DummySparseCategoricalCrossentropyLoss, self).__init__( losses.sparse_categorical_crossentropy, name=name, reduction=reduction, from_logits=from_logits, ) x = input_layer.Input(shape=[3]) y = core.Dense(10)(x) model = training.Model(x, y) model.compile( loss=DummySparseCategoricalCrossentropyLoss(from_logits=True)) model_round_trip = json.loads( json.dumps(model.loss, default=serialization.get_json_type)) # check if class name with package scope self.assertEqual("dummy-package>DummySparseCategoricalCrossentropyLoss", model_round_trip["class_name"]) # check if configure is correctly self.assertEqual(True, model_round_trip["config"]["from_logits"])
def testQuantize_RaisesErrorIfNoQuantizeProvider(self): annotated_model = keras.Sequential( [QuantizeAnnotate(self.CustomLayer(3), input_shape=(2, ))]) with generic_utils.custom_object_scope( {'CustomLayer': self.CustomLayer}): with self.assertRaises(RuntimeError): quantize_apply(annotated_model)
def testQuantize_PreferenceToUserSpecifiedQuantizeProvider(self): annotated_model = keras.Sequential([ QuantizeAnnotate(keras.layers.Dense(3), input_shape=(2, ), quantize_provider=_TestQuantizeProvider()) ]) with generic_utils.custom_object_scope( {'_TestQuantizeProvider': _TestQuantizeProvider}): quantized_model = quantize_apply(annotated_model) quantized_layer = quantized_model.layers[0] self._assert_layer_quantized(annotated_model.layers[0], quantized_layer) self.assertIsInstance(quantized_layer.quantize_provider, _TestQuantizeProvider)
def test_serializing_model_with_metric_with_custom_object_scope( self, value): def get_instance(x): if isinstance(x, str): return x if isinstance(x, type) and issubclass(x, metrics.Metric): return x() return x metric_input = nest.map_structure(get_instance, value) weighted_metric_input = nest.map_structure(get_instance, value) with generic_utils.custom_object_scope({ 'MyMeanAbsoluteError': MyMeanAbsoluteError, '_my_mae': _my_mae, 'Bias': testing_utils.Bias, }): model = _get_multi_io_model() model.compile(optimizer_v2.gradient_descent.SGD(0.1), 'mae', metrics=metric_input, weighted_metrics=weighted_metric_input, run_eagerly=testing_utils.should_run_eagerly(), experimental_run_tf_function=testing_utils. should_run_tf_function()) history = model.fit([self.x, self.x], [self.y, self.y], batch_size=3, epochs=3, sample_weight=[self.w, self.w]) # Assert training. self.assertAllClose(history.history['loss'], [2., 1.6, 1.2], 1e-3) eval_results = model.evaluate([self.x, self.x], [self.y, self.y], sample_weight=[self.w, self.w]) if h5py is None: return model.save(self.model_filename) loaded_model = keras.models.load_model(self.model_filename) loaded_model.predict([self.x, self.x]) loaded_eval_results = loaded_model.evaluate( [self.x, self.x], [self.y, self.y], sample_weight=[self.w, self.w]) # Assert all evaluation results are the same. self.assertAllClose(eval_results, loaded_eval_results, 1e-9)
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 testSerializationQuantizeAnnotate(self): input_shape = (2, ) layer = keras.layers.Dense(3) wrapper = quantize_annotate.QuantizeAnnotate( layer=layer, quantize_provider=self.TestQuantizeProvider(), input_shape=input_shape) custom_objects = { 'QuantizeAnnotate': quantize_annotate.QuantizeAnnotate, 'TestQuantizeProvider': self.TestQuantizeProvider } serialized_wrapper = serialize_layer(wrapper) with generic_utils.custom_object_scope(custom_objects): wrapper_from_config = deserialize_layer(serialized_wrapper) self.assertEqual(wrapper_from_config.get_config(), wrapper.get_config())
def testQuantize_UsesQuantizeProviderFromUser_NoBuiltIn(self): annotated_model = keras.Sequential([ QuantizeAnnotate(self.CustomLayer(3), input_shape=(2, ), quantize_provider=_TestQuantizeProvider()) ]) with generic_utils.custom_object_scope({ 'CustomLayer': self.CustomLayer, '_TestQuantizeProvider': _TestQuantizeProvider }): quantized_model = quantize_apply(annotated_model) quantized_layer = quantized_model.layers[0] self._assert_layer_quantized(annotated_model.layers[0], quantized_layer) self.assertIsInstance(quantized_layer.quantize_provider, _TestQuantizeProvider)
def prune_scope(): """Provides a scope in which Pruned layers and models can be deserialized. If a keras model or layer has been pruned, it needs to be within this scope to be successfully deserialized. Returns: Object of type `CustomObjectScope` with pruning objects included. Example: ```python pruned_model = prune_low_magnitude(model, **self.params) keras.models.save_model(pruned_model, keras_file) with prune_scope(): loaded_model = keras.models.load_model(keras_file) ``` """ return custom_object_scope( {'PruneLowMagnitude': pruning_wrapper.PruneLowMagnitude})
def testSerializationQuantizeWrapper(self): input_shape = (2, ) layer = keras.layers.Dense(3) wrapper = QuantizeWrapper( layer=layer, quantize_provider=self.quantize_registry.get_quantize_provider( layer), input_shape=input_shape) custom_objects = { 'QuantizeAwareActivation': QuantizeAwareActivation, 'QuantizeWrapper': QuantizeWrapper } custom_objects.update(tflite_quantize_registry._types_dict()) serialized_wrapper = serialize_layer(wrapper) with generic_utils.custom_object_scope(custom_objects): wrapper_from_config = deserialize_layer(serialized_wrapper) self.assertEqual(wrapper_from_config.get_config(), wrapper.get_config())
def cluster_scope(): """Provides a scope in which Clustered layers and models can be deserialized. If a keras model or layer has been clustered, it needs to be within this scope to be successfully deserialized. Returns: Object of type `CustomObjectScope` with clustering objects included. Example: ```python clustered_model = cluster_weights(model, **self.params) keras.models.save_model(clustered_model, keras_file) with cluster_scope(): loaded_model = keras.models.load_model(keras_file) ``` """ return custom_object_scope( {'ClusterWeights': cluster_wrapper.ClusterWeights})
def get_activation(identifier): if identifier is None: return None with custom_object_scope({'exp_relu': exp_relu}): return tf.keras.activations.get(identifier)