From 726de81edf6fb7d753d18057226754ecf13fc82a Mon Sep 17 00:00:00 2001 From: Ricardo Torres Date: Thu, 12 Nov 2026 09:27:14 -0800 Subject: [PATCH] Enforce Keras safe mode when loading models in evaluate_pruning. The evaluate_pruning tool now loads models with Keras safe mode enabled for all formats, including HDF5. Models with unsafe objects such as Lambda layers are rejected by default. Pass the new --allow_unsafe_deserialization flag to load models from trusted sources. PiperOrigin-RevId: 1020190437 --- .../python/core/sparsity/keras/tools/BUILD | 29 +++++- .../sparsity/keras/tools/evaluate_pruning.py | 64 ++++++++++++- .../keras/tools/evaluate_pruning_test.py | 93 +++++++++++++++++++ 3 files changed, 181 insertions(+), 5 deletions(-) create mode 100644 tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning_test.py diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/tools/BUILD b/tensorflow_model_optimization/python/core/sparsity/keras/tools/BUILD index 22d3d329e..5adc56afb 100644 --- a/tensorflow_model_optimization/python/core/sparsity/keras/tools/BUILD +++ b/tensorflow_model_optimization/python/core/sparsity/keras/tools/BUILD @@ -42,16 +42,41 @@ py_test( ], ) +py_library( + name = "evaluate_pruning_lib", + srcs = ["evaluate_pruning.py"], + # strict_deps = False, + deps = [ + ":sparsity_tooling", + # absl:app dep1, + # absl/flags dep1, + # tensorflow dep1, + "//tensorflow_model_optimization/python/core/keras:compat", + "//tensorflow_model_optimization/python/core/sparsity/keras:prune", + ], +) + py_binary( name = "evaluate_pruning", srcs = ["evaluate_pruning.py"], # strict_deps = False, deps = [ - ":sparsity_tooling", + ":evaluate_pruning_lib", # google/protobuf:use_fast_cpp_protos dep1, # Automatically added + ], +) + +py_test( + name = "evaluate_pruning_test", + size = "medium", + srcs = ["evaluate_pruning_test.py"], + deps = [ + ":evaluate_pruning_lib", + # absl/testing:parameterized dep1, + # google/protobuf:use_fast_cpp_protos dep1, + # numpy dep1, # tensorflow dep1, "//tensorflow_model_optimization/python/core/keras:compat", - "//tensorflow_model_optimization/python/core/sparsity/keras:prune", ], ) diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning.py b/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning.py index ab897d98c..a87f5a92b 100644 --- a/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning.py +++ b/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning.py @@ -25,6 +25,7 @@ from __future__ import print_function +import importlib import os import tempfile import textwrap @@ -51,6 +52,13 @@ 'block_size', '1,1', 'Comma-separated dimensions (height,weight) of the block sparsity pattern.' ) +_ALLOW_UNSAFE_DESERIALIZATION = flags.DEFINE_bool( + 'allow_unsafe_deserialization', + False, + 'Allow deserialization of Python bytecode embedded in the model (e.g. ' + '`Lambda` layers). This can lead to arbitrary code execution, so only ' + 'enable it for models from trusted sources.', +) def _parse_block_size_flag(value): @@ -67,6 +75,44 @@ def _check_block_size(flag_value): raise flags.ValidationError('Invalid block size value "%s".' % flag_value) +def _get_safe_mode_scope_class(): + """Returns Keras' `SafeModeScope` class.""" + # `SafeModeScope` lives next to `load_model` in the Keras `saving` package, + # whose import path differs between Keras distributions (e.g. + # `tf_keras.saving` vs `tf_keras.src.saving`). + saving_package = keras.models.load_model.__module__.rpartition('.')[0] + serialization_lib = importlib.import_module( + saving_package + '.serialization_lib' + ) + return serialization_lib.SafeModeScope + + +def load_model(model_path, allow_unsafe_deserialization=False): + """Loads a Keras model, disallowing unsafe deserialization by default. + + `keras.models.load_model(..., safe_mode=True)` only guards the native + `.keras` format. Legacy formats (e.g. HDF5) can embed Python bytecode in + `Lambda` layers, which would be executed on load. To protect against + arbitrary code execution, safe mode is enforced for every format through + Keras' `SafeModeScope`. + + Args: + model_path: Path to the Keras model to load. + allow_unsafe_deserialization: Whether to allow deserialization of Python + bytecode embedded in the model. Only set this for trusted models. + + Returns: + The loaded Keras model. + + Raises: + ValueError: If the model contains unsafe objects and + `allow_unsafe_deserialization` is False. + """ + safe_mode = not allow_unsafe_deserialization + with _get_safe_mode_scope_class()(safe_mode): + return keras.models.load_model(model_path, safe_mode=safe_mode) + + def convert_to_tflite(keras_model, output_path): """Converts the given Keras model to TFLite and write it to a file.""" converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) @@ -93,7 +139,13 @@ def pruned_model_filename(sparsity, block_size): sparsity, '%dx%d' % block_size) -def run(input_model_path, output_dir, target_sparsity, block_size): +def run( + input_model_path, + output_dir, + target_sparsity, + block_size, + allow_unsafe_deserialization=False, +): """Prunes the model and converts both pruned and unpruned versions to TFLite.""" print(textwrap.dedent("""\ @@ -101,7 +153,7 @@ def run(input_model_path, output_dir, target_sparsity, block_size): are not intended to be served in production, but to be used for performance benchmarking.""")) - input_model = keras.models.load_model(input_model_path) + input_model = load_model(input_model_path, allow_unsafe_deserialization) os.makedirs(output_dir, exist_ok=True) unpruned_tflite_path = os.path.join( @@ -136,7 +188,13 @@ def main(argv): raise app.UsageError('Too many command-line arguments.') block_size = _parse_block_size_flag(_BLOCK_SIZE.value) - run(_MODEL_PATH.value, _OUTPUT_DIR.value, _SPARSITY.value, block_size) + run( + _MODEL_PATH.value, + _OUTPUT_DIR.value, + _SPARSITY.value, + block_size, + _ALLOW_UNSAFE_DESERIALIZATION.value, + ) if __name__ == '__main__': diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning_test.py b/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning_test.py new file mode 100644 index 000000000..c82d6881b --- /dev/null +++ b/tensorflow_model_optimization/python/core/sparsity/keras/tools/evaluate_pruning_test.py @@ -0,0 +1,93 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for evaluate_pruning.py.""" + +import os + +from absl.testing import parameterized +import numpy as np +import tensorflow as tf + +from tensorflow_model_optimization.python.core.keras.compat import keras +from tensorflow_model_optimization.python.core.sparsity.keras.tools import evaluate_pruning + + +def _build_model(with_lambda): + layers = [keras.layers.Dense(4, input_shape=(3,))] + if with_lambda: + layers.append(keras.layers.Lambda(lambda x: x * 2.0)) + layers.append(keras.layers.Dense(2)) + return keras.Sequential(layers) + + +class EvaluatePruningLoadModelTest(tf.test.TestCase, parameterized.TestCase): + + def _save_model(self, model, filename): + path = os.path.join(self.create_tempdir().full_path, filename) + model.save(path) + return path + + @parameterized.parameters('model.h5', 'model.keras') + def test_load_model_rejects_lambda_by_default(self, filename): + path = self._save_model(_build_model(with_lambda=True), filename) + + with self.assertRaisesRegex(ValueError, 'Lambda'): + evaluate_pruning.load_model(path) + + @parameterized.parameters('model.h5', 'model.keras') + def test_load_model_allows_lambda_when_unsafe_allowed(self, filename): + model = _build_model(with_lambda=True) + path = self._save_model(model, filename) + + loaded_model = evaluate_pruning.load_model( + path, allow_unsafe_deserialization=True + ) + + x = np.ones((1, 3), dtype=np.float32) + self.assertAllClose(model.predict(x), loaded_model.predict(x)) + + @parameterized.parameters('model.h5', 'model.keras') + def test_load_model_loads_safe_model(self, filename): + model = _build_model(with_lambda=False) + path = self._save_model(model, filename) + + loaded_model = evaluate_pruning.load_model(path) + + x = np.ones((1, 3), dtype=np.float32) + self.assertAllClose(model.predict(x), loaded_model.predict(x)) + + def test_load_model_restores_safe_mode_state(self): + unsafe_path = self._save_model(_build_model(with_lambda=True), 'model.h5') + + evaluate_pruning.load_model(unsafe_path, allow_unsafe_deserialization=True) + + # Allowing unsafe deserialization once must not leak to later loads. + with self.assertRaisesRegex(ValueError, 'Lambda'): + evaluate_pruning.load_model(unsafe_path) + + def test_run_rejects_lambda_model(self): + path = self._save_model(_build_model(with_lambda=True), 'model.h5') + output_dir = self.create_tempdir().full_path + + with self.assertRaisesRegex(ValueError, 'Lambda'): + evaluate_pruning.run( + path, output_dir, target_sparsity=0.5, block_size=(1, 1) + ) + + self.assertEmpty(os.listdir(output_dir)) + + +if __name__ == '__main__': + tf.test.main()