def testConvertKerasModelToTfGraphModel(self):
output_dir = os.path.join(self._tmp_dir, 'foo_model')
sequential_model = keras.models.Sequential([
keras.layers.Dense(
3, input_shape=(2,), use_bias=True, kernel_initializer='ones',
name='Dense1')])
h5_path = os.path.join(self._tmp_dir, 'SequentialModel.h5')
sequential_model.save(h5_path)
converter.dispatch_keras_h5_to_tfjs_graph_model_conversion(
h5_path, output_dir=output_dir)
# Check model.json and weights manifest.
with open(os.path.join(output_dir, 'model.json'), 'rt') as f:
model_json = json.load(f)
self.assertTrue(model_json['modelTopology'])
weights_manifest = model_json['weightsManifest']
self.assertEqual(len(weights_manifest), 1)
# Check meta-data in the artifact JSON.
self.assertEqual(model_json['format'], 'graph-model')
self.assertEqual(
model_json['convertedBy'],
'TensorFlow.js Converter v%s' % version.version)
self.assertEqual(model_json['generatedBy'],
tf.__version__)
self.assertTrue(glob.glob(os.path.join(output_dir, 'group*-*')))
def testConvertKerasModelToTfGraphModelSharded(self):
output_dir = os.path.join(self._tmp_dir, 'foo_model')
sequential_model = tf.keras.models.Sequential([
tf.keras.layers.Dense(3,
input_shape=(2, ),
use_bias=True,
kernel_initializer='ones',
name='Dense1')
])
h5_path = os.path.join(self._tmp_dir, 'SequentialModel.h5')
sequential_model.save(h5_path)
# Do initial conversion without sharding.
converter.dispatch_keras_h5_to_tfjs_graph_model_conversion(
h5_path, output_dir)
weight_files = glob.glob(os.path.join(output_dir, 'group*.bin'))
# Get size of weights in bytes after graph optimizations.
optimized_total_weight = sum(
[os.path.getsize(f) for f in weight_files])
# Due to the shard size, there ought to be 4 shards after conversion.
weight_shard_size_bytes = int(optimized_total_weight * 0.3)
output_dir = os.path.join(self._tmp_dir, 'sharded_model')
# Convert Keras model again with shard argument set.
converter.dispatch_keras_h5_to_tfjs_graph_model_conversion(
h5_path,
output_dir,
weight_shard_size_bytes=weight_shard_size_bytes)
weight_files = sorted(glob.glob(os.path.join(output_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), optimized_total_weight)
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])