Add helper methods to load saved model from external files in model maker.
PiperOrigin-RevId: 480444918
This commit is contained in:
parent
12f72f067d
commit
f4fd1063a7
|
@ -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:
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue
Block a user