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)
