Add helper methods to load saved model from external files in model maker.

PiperOrigin-RevId: 480444918
This commit is contained in:
MediaPipe Team 2022-10-11 13:52:43 -07:00 committed by Copybara-Service
parent 12f72f067d
commit f4fd1063a7
3 changed files with 59 additions and 4 deletions

View File

@ -26,6 +26,7 @@ from typing import Any, Callable, Dict, List, Optional, Text, Tuple, Union
import numpy as np
import tensorflow as tf
# resources dependency
from mediapipe.model_maker.python.core.data import dataset
from mediapipe.model_maker.python.core.utils import quantization
@ -33,6 +34,31 @@ DEFAULT_SCALE, DEFAULT_ZERO_POINT = 0, 0
ESTIMITED_STEPS_PER_EPOCH = 1000
def load_keras_model(model_path: str,
compile_on_load: bool = False) -> tf.keras.Model:
"""Loads a tensorflow Keras model from file and returns the Keras model.
Args:
model_path: Relative path to a directory containing model data, such as
<parent_path>/saved_model/.
compile_on_load: Whether the model should be compiled while loading. If
False, the model returned has to be compiled with the appropriate loss
function and custom metrics before running for inference on a test
dataset.
Returns:
A tensorflow Keras model.
"""
# Extract the file path before mediapipe/ as the `base_dir`. By joining it
# with the `model_path` which defines the relative path under mediapipe/, it
# yields to the aboslution path of the model files directory.
cwd = os.path.dirname(__file__)
base_dir = cwd[:cwd.rfind('mediapipe')]
absolute_path = os.path.join(base_dir, model_path)
return tf.keras.models.load_model(
absolute_path, custom_objects={'tf': tf}, compile=compile_on_load)
def get_steps_per_epoch(steps_per_epoch: Optional[int] = None,
batch_size: Optional[int] = None,
train_data: Optional[dataset.Dataset] = None) -> int:

View File

@ -15,7 +15,6 @@
import os
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from mediapipe.model_maker.python.core.utils import model_util
@ -25,6 +24,18 @@ from mediapipe.model_maker.python.core.utils import test_util
class ModelUtilTest(tf.test.TestCase, parameterized.TestCase):
def test_load_model(self):
input_dim = 4
model = test_util.build_model(input_shape=[input_dim], num_classes=2)
saved_model_path = os.path.join(self.get_temp_dir(), 'saved_model')
model.save(saved_model_path)
loaded_model = model_util.load_keras_model(saved_model_path)
input_tensors = test_util.create_random_sample(size=[1, input_dim])
model_output = model.predict_on_batch(input_tensors)
loaded_model_output = loaded_model.predict_on_batch(input_tensors)
self.assertTrue((model_output == loaded_model_output).all())
@parameterized.named_parameters(
dict(
testcase_name='input_only_steps_per_epoch',
@ -124,9 +135,9 @@ class ModelUtilTest(tf.test.TestCase, parameterized.TestCase):
input_dim: int,
max_input_value: int = 1000,
atol: float = 1e-04):
np.random.seed(0)
random_input = np.random.uniform(
low=0, high=max_input_value, size=(1, input_dim)).astype(np.float32)
random_input = test_util.create_random_sample(
size=[1, input_dim], high=max_input_value)
random_input = tf.convert_to_tensor(random_input)
self.assertTrue(
test_util.is_same_output(

View File

@ -46,6 +46,24 @@ def create_dataset(data_size: int,
return dataset
def create_random_sample(size: Union[int, List[int]],
low: float = 0,
high: float = 1) -> np.ndarray:
"""Creates and returns a random sample with floating point values.
Args:
size: Size of the output multi-dimensional array.
low: Lower boundary of the output values.
high: Higher boundary of the output values.
Returns:
1D array if the size is scalar. Otherwise, N-D array whose dimension equals
input size.
"""
np.random.seed(0)
return np.random.uniform(low=low, high=high, size=size).astype(np.float32)
def build_model(input_shape: List[int], num_classes: int) -> tf.keras.Model:
"""Builds a simple Keras model for test."""
inputs = tf.keras.layers.Input(shape=input_shape)