def test_graph_models_to_saved_model(self):
"""graph_models_to_saved_model should accept model list"""
model_dir_1 = testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME)
model_dir_2 = testutils.get_path_to(testutils.PRELU_MODEL_PATH)
tags_1 = [tf.saved_model.SERVING, 'model_1']
tags_2 = [tf.saved_model.SERVING, 'model_2']
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
api.graph_models_to_saved_model([(model_dir_1, tags_1),
(model_dir_2, tags_2)],
export_dir)
self.assertTrue(os.path.exists(export_dir))
# try to load model 1
meta_graph_def = load_meta_graph(export_dir, tags_1)
self.assertIsNotNone(meta_graph_def)
# we also want a signature to be present
self.assertEqual(len(meta_graph_def.signature_def), 1)
# the signature should be valid
self.assertTrue(
is_valid_signature(
list(meta_graph_def.signature_def.values())[0]))
# try to load model 2
meta_graph_def = load_meta_graph(export_dir, tags_2)
self.assertIsNotNone(meta_graph_def)
# we also want a signature to be present
self.assertEqual(len(meta_graph_def.signature_def), 1)
# the signature should be valid
self.assertTrue(
is_valid_signature(
list(meta_graph_def.signature_def.values())[0]))
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_graph_models_to_saved_model_accepts_signatures(self):
"""graph_models_to_saved_model should accept signatures"""
model_dir_1 = testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME)
model_dir_2 = testutils.get_path_to(testutils.MULTI_HEAD_PATH)
tags_1 = [tf.saved_model.SERVING, 'model_1']
tags_2 = [tf.saved_model.SERVING, 'model_2']
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
model_list = [(model_dir_1, tags_1), (model_dir_2, tags_2)]
signatures = {
'ignore_this': {'': {api.SIGNATURE_OUTPUTS: ['y']}},
model_dir_2: {'': {api.SIGNATURE_OUTPUTS: ['Identity']}}
}
api.graph_models_to_saved_model(model_list, export_dir, signatures)
self.assertTrue(os.path.exists(export_dir))
# try to load model 2
meta_graph_def = load_meta_graph(export_dir, tags_2)
# we want a signature to be present
self.assertEqual(len(meta_graph_def.signature_def), 1)
signature = list(meta_graph_def.signature_def.values())[0]
# the signature should be valid
self.assertTrue(is_valid_signature(signature))
# the signature should have a single output
self.assertEqual(len(signature.outputs), 1)
self.assertIn('Identity', signature.outputs.keys())
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_graph_model_to_saved_model_accepts_signature_key_map(self):
"""graph_model_to_saved_model should accept signature key map"""
model_dir = testutils.get_path_to(testutils.MULTI_HEAD_PATH)
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
tags = [tf.saved_model.SERVING]
signature_map = {
'': {api.SIGNATURE_OUTPUTS: ['Identity']},
'debug': {api.SIGNATURE_OUTPUTS: ['Identity', 'Identity_1']}}
signature_key = api.RenameMap([
('Identity', 'output'), ('Identity_1', 'autoencoder_output')
])
api.graph_model_to_saved_model(model_dir, export_dir,
tags=tags,
signature_def_map=signature_map,
signature_key_map=signature_key)
# try and load the model
meta_graph_def = load_meta_graph(export_dir, tags)
# the signatures should contain the renamed keys
for signature in meta_graph_def.signature_def.values():
self.assertIn('output', signature.outputs)
self.assertEqual(signature.outputs['output'].name,
'Identity:0')
signature = meta_graph_def.signature_def['debug']
self.assertIn('autoencoder_output', signature.outputs)
self.assertEqual(signature.outputs['autoencoder_output'].name,
'Identity_1:0')
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_graph_model_to_saved_model_accepts_signature_map(self):
"""graph_model_to_saved_model should accept signature map"""
model_dir = testutils.get_path_to(testutils.MULTI_HEAD_PATH)
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
tags = [tf.saved_model.SERVING]
signature_map = {
'': {api.SIGNATURE_OUTPUTS: ['Identity']},
'debug': {api.SIGNATURE_OUTPUTS: ['Identity', 'Identity_1']}}
api.graph_model_to_saved_model(model_dir, export_dir,
tags=tags,
signature_def_map=signature_map)
# try and load the model
meta_graph_def = load_meta_graph(export_dir, tags)
self.assertIsNotNone(meta_graph_def)
# we want both signatures to be present
self.assertEqual(len(meta_graph_def.signature_def), 2)
# the signatures should be valid
for signature in meta_graph_def.signature_def.values():
self.assertTrue(is_valid_signature(signature))
# the default signature should have one output
default_key = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
self.assertEqual(
len(meta_graph_def.signature_def[default_key].outputs), 1)
# debug signature should be present and contain two outputs
self.assertIn('debug', meta_graph_def.signature_def.keys())
self.assertEqual(
len(meta_graph_def.signature_def['debug'].outputs), 2)
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def main(args: List[str]) -> None:
if len(args) != 2:
print(f"usage: {os.path.basename(args[0])} IMAGE")
exit(1)
if not os.path.isfile(args[1]):
print(f"{os.path.basename(args[0])}: '{args[1]}' isn't a file")
exit(1)
tfjs.api.enable_cuda()
image = tf.keras.preprocessing.image.load_img(args[1])
image = tf.keras.preprocessing.image.img_to_array(image)
# normalise image data to [-1, 1]
image /= 127.5
image -= 1.
# ensure image size matches model input shape
if image.shape != (32, 32, 3):
print(f"{os.path.basename(args[0])}: WARNING - image size "
f"should be 32x32, not {image.shape[0]}x{image.shape[1]}")
image = Resizing(height=32, width=32)(image)
# reshape to fit model input (and convert to tensor if necessary)
image = tf.reshape(image, [1, 32, 32, 3])
# grab the model file and convert graph to function
graph = util.get_sample_graph(util.get_path_to(util.DEPTHWISE_RELU_FILE))
model = tfjs.api.graph_to_function_v2(graph)
# run the model on the input image
result = model(image)
# show result
label, confidence = _evaluate(result)
print(f"Result: {label}, confidence={confidence}")
def test_graph_models_to_saved_model_accepts_signature_keys(self):
"""graph_models_to_saved_model should accept signature keys"""
model_dir_1 = testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME)
model_dir_2 = testutils.get_path_to(testutils.MULTI_HEAD_PATH)
tags_1 = [tf.saved_model.SERVING, 'model_1']
tags_2 = [tf.saved_model.SERVING, 'model_2']
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
model_list = [(model_dir_1, tags_1), (model_dir_2, tags_2)]
signatures = {
'ignore_this': {
'': {
api.SIGNATURE_OUTPUTS: ['y']
}
},
model_dir_2: {
'': {
api.SIGNATURE_OUTPUTS: ['Identity']
}
}
}
signature_keys = {
model_dir_1: api.RenameMap({
'x': 'input',
'Identity': 'output'
}),
model_dir_2: api.RenameMap({'Identity': 'scores'})
}
api.graph_models_to_saved_model(model_list, export_dir, signatures,
signature_keys)
self.assertTrue(os.path.exists(export_dir))
# check the signatures of model 1
meta_graph_def = load_meta_graph(export_dir, tags_1)
signature = list(meta_graph_def.signature_def.values())[0]
self.assertIn('input', signature.inputs)
self.assertEqual(signature.inputs['input'].name, 'x:0')
self.assertIn('output', signature.outputs)
self.assertEqual(signature.outputs['output'].name, 'Identity:0')
# check the signatures of model 2
meta_graph_def = load_meta_graph(export_dir, tags_2)
signature = list(meta_graph_def.signature_def.values())[0]
self.assertIn('scores', signature.outputs)
self.assertEqual(signature.outputs['scores'].name, 'Identity:0')
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_load_graph_model_with_fused_depthwise_prelu(self):
"""load_graph_model should split fused depthwise conv2d with prelu"""
model_dir = testutils.get_path_to(testutils.DEPTHWISE_PRELU_PATH)
graph = api.load_graph_model(model_dir)
loaded_model = testutils.graph_to_model(graph)
original_model_name = testutils.get_path_to(
testutils.DEPTHWISE_PRELU_FILE)
original_model = testutils.graph_to_model(original_model_name)
# run both models and compare results
x = _load_image(testutils.get_path_to('./data/human1.jpg'))
y_from_loaded = _argmax(loaded_model(x))
y_from_original = _argmax(original_model(x))
# sanity check - should be reckognised as a human (class 1)
self.assertEqual(y_from_original[0], 1)
# same class
self.assertEqual(y_from_loaded[0], y_from_original[0])
# same confidence
self.assertAlmostEqual(y_from_loaded[1], y_from_original[1], 4)
def test_load_graph_model_with_simple_model(self):
"""load_graph_model should load simple model"""
model_dir = testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME)
graph = api.load_graph_model(model_dir)
self.assertIsInstance(graph, tf.Graph)
loaded_model = testutils.graph_to_model(graph)
original_model_name = testutils.get_path_to(
testutils.SIMPLE_MODEL_FILE_NAME)
original_model = testutils.graph_to_model(original_model_name)
# run both models and compare results
x_ = 4
x = tf.constant([[x_]], dtype=tf.float32)
y_from_loaded_model = as_scalar(loaded_model(x))
y_from_original_model = as_scalar(original_model(x))
# sanity check; fails if model is different from the one we expected:
# we want a model that predicts y = 5*x
self.assertAlmostEqual(y_from_original_model, x_*5, places=1)
# actual test
self.assertAlmostEqual(y_from_loaded_model, y_from_original_model,
places=4)
def test_graph_models_to_saved_model(self):
"""graph_models_to_saved_model should accept model list"""
model_dir_1 = testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME)
model_dir_2 = testutils.get_path_to(testutils.PRELU_MODEL_PATH)
tags_1 = ['serving_default', 'model_1']
tags_2 = ['serving_default', 'model_2']
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
api.graph_models_to_saved_model([(model_dir_1, tags_1),
(model_dir_2, tags_2)],
export_dir)
self.assertTrue(os.path.exists(export_dir))
# try to load model 1
model = tf.saved_model.load(export_dir, tags=tags_1)
self.assertIsNotNone(model.graph)
# try toload model 2
model = tf.saved_model.load(export_dir, tags=tags_2)
self.assertIsNotNone(model.graph)
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_load_graph_model_with_prelu(self):
"""load_graph_model should convert prelu operations"""
model_dir = testutils.get_path_to(testutils.PRELU_MODEL_PATH)
graph = api.load_graph_model(model_dir)
loaded_model = testutils.graph_to_model(graph)
original_model_name = testutils.get_path_to(testutils.PRELU_MODEL_FILE)
original_model = testutils.graph_to_model(original_model_name)
# run both models and compare results
cx, cy, cz, r = -0.12, 0.2, 0.1, 0.314158
px, py, pz = -0.4, 0.5, 0.4
x = tf.constant([[cx, cy, cz, r, px, py, pz]], dtype=tf.float32)
y_from_loaded_model = as_scalar(loaded_model(x))
y_from_original_model = as_scalar(original_model(x))
# sanity check; fails if model is different from the one we expected:
# we want a model that predicts whether a point (px,py,pz) is inside
# a sphere at (cx,cy,cz) of radius r
self.assertAlmostEqual(y_from_original_model, 1, places=1)
# actual test
self.assertAlmostEqual(y_from_loaded_model,
y_from_original_model,
places=4)
def test_rename_map_apply(self):
"""RenameMap.apply maps old names to new names"""
_, signature_def = api.load_graph_model_and_signature(
testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME))
mapping = api.RenameMap({'x': 'input', 'Identity': 'output'})
updated = mapping.apply(signature_def)
self.assertNotIn('x', updated.inputs)
self.assertIn('input', updated.inputs)
# keep the tensor name!
self.assertEqual(updated.inputs['input'].name, 'x:0')
self.assertNotIn('Identity', updated.outputs)
self.assertIn('output', updated.outputs)
# keep the tensor name!
self.assertEqual(updated.outputs['output'].name, 'Identity:0')
def _load_hoh_dataset(tmpdirname):
# load dataset from archive; return dataset tuple for training and testing
archive = get_path_to(IMAGE_DATASET)
with ZipFile(archive, 'r') as zip:
zip.extractall(tmpdirname)
path = os.path.join(tmpdirname, 'train')
norm = Rescaling(1./127.5, offset=-1)
train_ds = image_dataset_from_directory(path, image_size=(32, 32), seed=42)
train_ds = train_ds.map(
lambda x, y: (norm(x), tf.one_hot(y, depth=2)))
path = os.path.join(tmpdirname, 'test')
val_ds = image_dataset_from_directory(path, image_size=(32, 32), seed=42)
val_ds = val_ds.map(
lambda x, y: (norm(x), tf.one_hot(y, depth=2)))
return (train_ds, val_ds)
def test_graph_model_to_saved_model(self):
"""graph_model_to_saved_model should save valid SavedModel"""
model_dir = testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME)
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
tags = ['serving_default']
api.graph_model_to_saved_model(model_dir, export_dir, tags=tags)
self.assertTrue(os.path.exists(export_dir))
# must be valid model; tf.saved_model.contains_saved_model is
# insufficient
imported = tf.saved_model.load(export_dir, tags=tags)
self.assertIsNotNone(imported.graph)
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_build_signatures(self):
"""_build_signatures should apply given key and include inputs"""
graph = testutils.get_sample_graph(
testutils.get_path_to(testutils.MULTI_HEAD_FILE))
default_key = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
debug_key = 'debug_model'
signature_map = {
'': {api.SIGNATURE_OUTPUTS: ['Identity']},
debug_key: {api.SIGNATURE_OUTPUTS: ['Identity', 'Identity_1']}}
signature_def_map = api._build_signatures(graph, signature_map)
self.assertIn(default_key, signature_def_map)
self.assertIn(debug_key, signature_def_map)
for signature_def in signature_def_map.values():
self.assertTrue(is_valid_signature(signature_def))
def test_graph_model_to_frozen_graph(self):
"""graph_model_to_frozen_graph should save valid frozen graph model"""
try:
input_name = testutils.get_path_to(
testutils.SIMPLE_MODEL_PATH_NAME)
output_name = os.path.join(tempfile.gettempdir(), 'frozen.pb')
api.graph_model_to_frozen_graph(input_name, output_name)
# make sure the output file exists and isn't empty
self.assertTrue(os.path.exists(output_name))
self.assertGreater(os.stat(output_name).st_size, 256)
# file must be a valid protobuf message
with open(output_name, 'rb') as pb_file:
graph_def = testutils.GraphDef()
graph_def.ParseFromString(pb_file.read())
finally:
if os.path.exists(output_name):
os.remove(output_name)
def test_load_graph_model_and_signature_from_meta_data(self):
"""load_graph_model_and_signature should extract signature def"""
_, signature_def = api.load_graph_model_and_signature(
testutils.get_path_to(testutils.PRELU_MODEL_PATH))
self.assertIsInstance(signature_def, util.SignatureDef)
self.assertTrue(is_valid_signature(signature_def))
self.assertEqual(len(signature_def.inputs), 1)
key, value = list(signature_def.inputs.items())[0]
self.assertEqual(key, 'input_vector')
self.assertEqual(value.name, 'input_vector:0')
self.assertEqual(value.dtype, tf.dtypes.float32)
self.assertEqual(_shape_of(value), (-1, 7))
self.assertEqual(len(signature_def.outputs), 1)
key, value = list(signature_def.outputs.items())[0]
self.assertEqual(key, 'Identity')
self.assertEqual(value.name, 'Identity:0')
self.assertEqual(value.dtype, tf.dtypes.float32)
self.assertEqual(_shape_of(value), (-1, 1))
def test_rename_input_nodes(self):
"""rename_input_nodes should rename input nodes in-place"""
model_file = testutils.get_path_to(testutils.SIMPLE_MODEL_FILE_NAME)
graph_def = testutils.get_sample_graph_def(model_file)
updated = util.rename_input_nodes(graph_def, {'x': 'scalar'})
# update should be in-place
self.assertEqual(graph_def, updated)
# inputs should be renamed
self.assertEqual(util.get_input_nodes(updated)[0].name, 'scalar')
# model should still work
model = testutils.graph_to_model(updated)
s = 18
scalar = tf.constant([[s]], dtype=tf.float32)
result = model(scalar)
value = result[0].numpy()
# value = np.reshape(value, (1))
y = value[0]
self.assertAlmostEqual(y, s*5, delta=0.1)
def test_rename_input_nodes_reject_invalid_args(self):
"""rename_input_nodes rejects invalid arguments"""
model_file = testutils.get_path_to(testutils.SIMPLE_MODEL_FILE_NAME)
graph_def = testutils.get_sample_graph_def(model_file)
# reject unknown node
self.assertRaises(
ValueError, lambda: util.rename_input_nodes(
graph_def, {'does-not-exist': 'scalar'}))
# reject non-input node
self.assertRaises(
ValueError, lambda: util.rename_input_nodes(
graph_def, {'Identity': 'scalar'}))
# reject rename to existing node
self.assertRaises(
ValueError, lambda: util.rename_input_nodes(
graph_def, {'x': 'Identity'}))
# new name must differ from old name
self.assertRaises(
ValueError, lambda: util.rename_input_nodes(
graph_def, {'x': 'x'}))
def test_load_graph_model_and_signature_from_tree(self):
"""load_graph_model_and_signature should infer signature def
from graph if signature def is incomplete
"""
_, signature_def = api.load_graph_model_and_signature(
testutils.get_path_to(testutils.SIMPLE_MODEL_PATH_NAME))
# simple model is missing inputs in signature - defer from graph
self.assertIsInstance(signature_def, util.SignatureDef)
self.assertTrue(is_valid_signature(signature_def))
self.assertEqual(len(signature_def.inputs), 1)
key, value = list(signature_def.inputs.items())[0]
self.assertEqual(key, 'x')
self.assertEqual(value.name, 'x:0')
self.assertEqual(value.dtype, tf.dtypes.float32)
self.assertEqual(_shape_of(value), (-1, 1))
self.assertEqual(len(signature_def.outputs), 1)
key, value = list(signature_def.outputs.items())[0]
self.assertEqual(key, 'Identity')
self.assertEqual(value.name, 'Identity:0')
self.assertEqual(value.dtype, tf.dtypes.float32)
self.assertEqual(_shape_of(value), (-1, 1))
def test_graph_model_to_saved_model(self):
"""graph_model_to_saved_model should save valid SavedModel"""
model_dir = testutils.get_path_to(testutils.PRELU_MODEL_PATH)
export_dir = tempfile.mkdtemp(suffix='.saved_model')
try:
tags = [tf.saved_model.SERVING]
api.graph_model_to_saved_model(model_dir, export_dir, tags=tags)
self.assertTrue(os.path.exists(export_dir))
self.assertTrue(tf.saved_model.contains_saved_model(export_dir))
# try and load the model
meta_graph_def = load_meta_graph(export_dir, tags)
self.assertIsNotNone(meta_graph_def)
# we also want a signature to be present
self.assertEqual(len(meta_graph_def.signature_def), 1)
# the signatures should be valid
self.assertTrue(
is_valid_signature(
list(meta_graph_def.signature_def.values())[0]))
finally:
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
def test_rename_output_nodes_append_identity(self):
"""rename_output_nodes should work for outputs that aren't Identity"""
model_file = testutils.get_path_to(testutils.SIMPLE_MODEL_FILE_NAME)
graph_def = testutils.get_sample_graph_def(model_file)
# some open-heart surgery on the model to remove the "Identity" output
idx = [i for (i, n) in enumerate(graph_def.node) if n.op == 'Identity']
del graph_def.node[idx[0]]
output = util.get_output_nodes(graph_def)[0].name
updated = util.rename_output_nodes(graph_def, {output: 'estimate'})
# update should be in-place
self.assertEqual(graph_def, updated)
# outputs should be renamed
self.assertEqual(util.get_output_nodes(updated)[0].name, 'estimate')
# model should still work
model = testutils.graph_to_model(updated)
s = 18
scalar = tf.constant([[s]], dtype=tf.float32)
result = model(scalar)
value = result[0].numpy()
# value = np.reshape(value, (1))
y = value[0]
self.assertAlmostEqual(y, s*5, delta=0.1)
def save_keras_model(model: Callable, path: str) -> None:
"""Save Keras model as TFJS graph model"""
Path(path).mkdir(parents=True, exist_ok=True)
model.save(os.path.join(path, 'keras.h5'))
convert_to_tfjs([
'--input_format=keras',
'--output_format=tfjs_graph_model',
os.path.join(path, 'keras.h5'), path
])
if __name__ == '__main__':
print('Generating multi-layer sample model...')
model = deepmind_atari_net(10, input_shape=(128, 128, 3))
save_tf_model(model, get_path_to(SAMPLE_MODEL_FILE_NAME))
print('Generating single-layer simple model...')
model = simple_model()
save_tfjs_model(model, get_path_to(SIMPLE_MODEL_PATH_NAME))
print('Generating prelu-activation model...')
model = prelu_classifier_model()
save_keras_model(model, get_path_to(PRELU_MODEL_PATH))
print('Generating multi-head model...')
model = multi_head_model()
save_tfjs_model(model, get_path_to(MULTI_HEAD_PATH))
print('Generating depthwise conv2d model...')
model = depthwise_model()
save_tfjs_model(model, get_path_to(DEPTHWISE_RELU_PATH))
print('Generating depthwise conv2d model with PReLU activation...')
model = depthwise_model('prelu')
save_tfjs_model(model, get_path_to(DEPTHWISE_PRELU_PATH))
def test_load_graph_model_and_signature_format_check(self):
"""load_graph_model_and_signature verifies model format"""
model = testutils.get_path_to(testutils.KERAS_MODEL_FILE_NAME)
self.assertRaises(api.ModelFormatError,
lambda: api.load_graph_model_and_signature(model))
