Enforce Keras safe mode when loading models in evaluate_pruning.

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",
     ],
 )
 

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,15 +139,21 @@ 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("""\
     Warning: The sparse models produced by this tool have poor accuracy. They
              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__':

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()