def testSavedModelSucceedsForCustomShardSize(self):
model = tf.keras.Sequential([
tf.keras.layers.Dense(1, input_shape=[2], activation='relu'),
tf.keras.layers.Dense(3, activation='tanh')
])
weights = model.get_weights()
total_weight_bytes = sum(np.size(w) for w in weights) * 4
# Due to the shard size, there ought to be 4 shards after conversion.
weight_shard_size_bytes = int(total_weight_bytes * 0.3)
# Convert Keras model to tfjs_layers_model format.
conversion.save_keras_model(
model,
self._tmp_dir,
weight_shard_size_bytes=weight_shard_size_bytes)
weight_files = sorted(
glob.glob(os.path.join(self._tmp_dir, 'group*.bin')))
self.assertEqual(len(weight_files), 4)
weight_file_sizes = [os.path.getsize(f) for f in weight_files]
self.assertEqual(sum(weight_file_sizes), total_weight_bytes)
self.assertEqual(weight_file_sizes[0], weight_file_sizes[1])
self.assertEqual(weight_file_sizes[0], weight_file_sizes[2])
self.assertLess(weight_file_sizes[3], weight_file_sizes[0])
def testConvertModelWithNestedLayerNames(self):
model = keras.Sequential()
# Add a layer with a nested layer name, i.e., a layer name with slash(es)
# in it.
model.add(keras.layers.Dense(2, input_shape=[12], name='dense'))
model.add(keras.layers.Dense(8, name='foo/dense'))
model.add(keras.layers.Dense(4, name='foo/bar/dense'))
tfjs_path = os.path.join(self._tmp_dir, 'nested_layer_names_model')
keras_h5_conversion.save_keras_model(model, tfjs_path)
# Check model.json and weights manifest.
with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f:
model_json = json.load(f)
self.assertTrue(model_json['modelTopology'])
weights_manifest = model_json['weightsManifest']
weight_shapes = dict()
for group in weights_manifest:
for weight in group['weights']:
weight_shapes[weight['name']] = weight['shape']
self.assertEqual(
sorted([
'dense/kernel', 'dense/bias', 'foo/dense/kernel',
'foo/dense/bias', 'foo/bar/dense/kernel', 'foo/bar/dense/bias'
]), sorted(list(weight_shapes.keys())))
self.assertEqual([12, 2], weight_shapes['dense/kernel'])
self.assertEqual([2], weight_shapes['dense/bias'])
self.assertEqual([2, 8], weight_shapes['foo/dense/kernel'])
self.assertEqual([8], weight_shapes['foo/dense/bias'])
self.assertEqual([8, 4], weight_shapes['foo/bar/dense/kernel'])
self.assertEqual([4], weight_shapes['foo/bar/dense/bias'])
def testSaveModelSucceedsForNestedKerasModel(self):
inner_model = keras.Sequential([
keras.layers.Dense(4, input_shape=[3], activation='relu'),
keras.layers.Dense(3, activation='tanh')
])
outer_model = keras.Sequential()
outer_model.add(inner_model)
outer_model.add(keras.layers.Dense(1, activation='sigmoid'))
keras_h5_conversion.save_keras_model(outer_model, self._tmp_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
# Verify that all the layers' weights are present.
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
weight_entries = []
for group in weights_manifest:
weight_entries.extend(group['weights'])
self.assertEqual(6, len(weight_entries))
def testSaveModelSucceedsForTfKerasNonSequentialModel(self):
t_input = tf.keras.Input([2])
dense_layer = tf.keras.layers.Dense(3)
t_output = dense_layer(t_input)
model = tf.keras.Model(t_input, t_output)
# `tf.keras.Model`s must be compiled before they can be saved.
model.compile(loss='mean_squared_error', optimizer='sgd')
keras_h5_conversion.save_keras_model(model, self._tmp_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
self.assertEqual(1, len(weights_manifest))
self.assertIn('paths', weights_manifest[0])
def testConvertModelWithNestedLayerNames(self):
model = keras.Sequential()
# Add a layer with a nested layer name, i.e., a layer name with slash(es)
# in it.
model.add(keras.layers.Dense(2, input_shape=[12], name='dense'))
model.add(keras.layers.Dense(8, name='foo/dense'))
model.add(keras.layers.Dense(4, name='foo/bar/dense'))
tfjs_path = os.path.join(self._tmp_dir, 'nested_layer_names_model')
keras_h5_conversion.save_keras_model(model, tfjs_path)
# Check model.json and weights manifest.
with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f:
model_json = json.load(f)
self.assertTrue(model_json['modelTopology'])
weights_manifest = model_json['weightsManifest']
weight_shapes = dict()
for group in weights_manifest:
for weight in group['weights']:
weight_shapes[weight['name']] = weight['shape']
self.assertEqual(
sorted(['dense/kernel', 'dense/bias', 'foo/dense/kernel',
'foo/dense/bias', 'foo/bar/dense/kernel',
'foo/bar/dense/bias']),
sorted(list(weight_shapes.keys())))
self.assertEqual([12, 2], weight_shapes['dense/kernel'])
self.assertEqual([2], weight_shapes['dense/bias'])
self.assertEqual([2, 8], weight_shapes['foo/dense/kernel'])
self.assertEqual([8], weight_shapes['foo/dense/bias'])
self.assertEqual([8, 4], weight_shapes['foo/bar/dense/kernel'])
self.assertEqual([4], weight_shapes['foo/bar/dense/bias'])
def testSaveModelSucceedsForNestedKerasModel(self):
inner_model = keras.Sequential([
keras.layers.Dense(4, input_shape=[3], activation='relu'),
keras.layers.Dense(3, activation='tanh')])
outer_model = keras.Sequential()
outer_model.add(inner_model)
outer_model.add(keras.layers.Dense(1, activation='sigmoid'))
keras_h5_conversion.save_keras_model(outer_model, self._tmp_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
# Verify that all the layers' weights are present.
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
weight_entries = []
for group in weights_manifest:
weight_entries.extend(group['weights'])
self.assertEqual(6, len(weight_entries))
def testSaveModelSucceedsForTfKerasNonSequentialModel(self):
t_input = tf.keras.Input([2])
dense_layer = tf.keras.layers.Dense(3)
t_output = dense_layer(t_input)
model = tf.keras.Model(t_input, t_output)
# `tf.keras.Model`s must be compiled before they can be saved.
model.compile(loss='mean_squared_error', optimizer='sgd')
keras_h5_conversion.save_keras_model(model, self._tmp_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
self.assertEqual(1, len(weights_manifest))
self.assertIn('paths', weights_manifest[0])
def testLoadFunctionalTfKerasModel(self):
with tf.Graph().as_default(), tf.compat.v1.Session():
input1 = keras.Input([4])
x1 = keras.layers.Dense(2, activation='relu')(input1)
x1 = keras.layers.BatchNormalization()(x1)
input2 = keras.Input([10])
x2 = keras.layers.Dense(5, activation='relu')(input2)
x2 = keras.layers.BatchNormalization()(x2)
y = keras.layers.Concatenate()([x1, x2])
y = keras.layers.Dense(1, activation='sigmoid')(y)
model = keras.Model([input1, input2], y)
model.compile(loss='binary_crossentropy', optimizer='sgd')
input1_val = np.ones([1, 4])
input2_val = np.ones([1, 10])
predict_out = model.predict([input1_val, input2_val])
save_dir = os.path.join(self._tmp_dir, 'functional_model')
keras_h5_conversion.save_keras_model(model, save_dir)
with tf.Graph().as_default(), tf.compat.v1.Session():
model2 = keras_tfjs_loader.load_keras_model(
os.path.join(save_dir, 'model.json'))
self.assertAllClose(predict_out,
model2.predict([input1_val, input2_val]))
def _saveKerasModelForTest(self, path):
model = keras.Sequential()
model.add(keras.layers.Dense(
2, input_shape=[12], bias_initializer='random_normal', name='dense'))
model.add(keras.layers.Dense(
8, bias_initializer='random_normal', name='foo/dense'))
model.add(keras.layers.Dense(
4, bias_initializer='random_normal', name='foo/bar/dense'))
keras_h5_conversion.save_keras_model(model, path)
return model
def testSavedModelRaisesErrorIfArtifactsDirExistsAsAFile(self):
artifacts_dir = os.path.join(self._tmp_dir, 'artifacts')
with open(artifacts_dir, 'wt') as f:
f.write('foo\n')
t_input = keras.Input([2])
dense_layer = keras.layers.Dense(3)
t_output = dense_layer(t_input)
model = keras.Model(t_input, t_output)
with self.assertRaisesRegexp(ValueError, r'already exists as a file'):
keras_h5_conversion.save_keras_model(model, artifacts_dir)
def testSavedModelRaisesErrorIfArtifactsDirExistsAsAFile(self):
artifacts_dir = os.path.join(self._tmp_dir, 'artifacts')
with open(artifacts_dir, 'wt') as f:
f.write('foo\n')
t_input = keras.Input([2])
dense_layer = keras.layers.Dense(3)
t_output = dense_layer(t_input)
model = keras.Model(t_input, t_output)
with self.assertRaisesRegexp( # pylint: disable=deprecated-method
ValueError, r'already exists as a file'):
keras_h5_conversion.save_keras_model(model, artifacts_dir)
def testLoadNestedKerasModel(self):
with tf.Graph().as_default(), tf.compat.v1.Session():
inner_model = keras.Sequential([
keras.layers.Dense(4, input_shape=[3], activation='relu'),
keras.layers.Dense(3, activation='tanh')])
outer_model = keras.Sequential()
outer_model.add(inner_model)
outer_model.add(keras.layers.Dense(1, activation='sigmoid'))
x = np.ones([1, 3], dtype=np.float32)
predict_out = outer_model.predict(x)
save_dir = os.path.join(self._tmp_dir, 'nested_model')
keras_h5_conversion.save_keras_model(outer_model, save_dir)
with tf.Graph().as_default(), tf.compat.v1.Session():
model2 = keras_tfjs_loader.load_keras_model(
os.path.join(save_dir, 'model.json'), use_unique_name_scope=True)
self.assertAllClose(predict_out, model2.predict(x))
def testSavedModelSucceedsForExistingDirAndSequential(self):
artifacts_dir = os.path.join(self._tmp_dir, 'artifacts')
os.makedirs(artifacts_dir)
model = keras.Sequential()
model.add(keras.layers.Dense(3, input_shape=[2]))
keras_h5_conversion.save_keras_model(model, artifacts_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(artifacts_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(artifacts_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
self.assertEqual(1, len(weights_manifest))
self.assertIn('paths', weights_manifest[0])
def testSavedModelSucceedsForExistingDirAndSequential(self):
artifacts_dir = os.path.join(self._tmp_dir, 'artifacts')
os.makedirs(artifacts_dir)
model = keras.Sequential()
model.add(keras.layers.Dense(3, input_shape=[2]))
keras_h5_conversion.save_keras_model(model, artifacts_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(artifacts_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(artifacts_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
self.assertEqual(1, len(weights_manifest))
self.assertIn('paths', weights_manifest[0])
def testSaveModelSucceedsForNonSequentialModel(self):
t_input = keras.Input([2])
dense_layer = keras.layers.Dense(3)
t_output = dense_layer(t_input)
model = keras.Model(t_input, t_output)
keras_h5_conversion.save_keras_model(model, self._tmp_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
self.assertEqual(1, len(weights_manifest))
self.assertIn('paths', weights_manifest[0])
def testSaveModelSucceedsForNonSequentialModel(self):
t_input = keras.Input([2])
dense_layer = keras.layers.Dense(3)
t_output = dense_layer(t_input)
model = keras.Model(t_input, t_output)
keras_h5_conversion.save_keras_model(model, self._tmp_dir)
# Verify the content of the artifacts output directory.
self.assertTrue(
os.path.isfile(os.path.join(self._tmp_dir, 'group1-shard1of1')))
model_json = json.load(
open(os.path.join(self._tmp_dir, 'model.json'), 'rt'))
topology_json = model_json['modelTopology']
self.assertIn('keras_version', topology_json)
self.assertIn('backend', topology_json)
self.assertIn('model_config', topology_json)
weights_manifest = model_json['weightsManifest']
self.assertTrue(isinstance(weights_manifest, list))
self.assertEqual(1, len(weights_manifest))
self.assertIn('paths', weights_manifest[0])
def _saveRNNKerasModelForTest(self, path):
model = tf.keras.Sequential()
model.add(tf.keras.layers.Embedding(100, 20, input_shape=[10]))
model.add(tf.keras.layers.SimpleRNN(4))
keras_h5_conversion.save_keras_model(model, path)
return model
