def test_optimizer_add_unsupported_op(self): self._create_unsupported_saved_model() with self.assertRaisesRegexp( # pylint: disable=deprecated-method ValueError, r'^Unsupported Ops'): tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR))
def test_convert_saved_model_v1(self): self._create_saved_model_v1() input_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) output_dir = os.path.join(input_dir, 'js') tf_saved_model_conversion_v2.convert_tf_saved_model( input_dir, output_dir ) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'js') # 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'] self.assertCountEqual(weights_manifest[0]['paths'], ['group1-shard1of1.bin']) self.assertIn('weights', weights_manifest[0]) # 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 test_convert_saved_model_strip_debug_ops(self): self._create_saved_model_with_debug_ops() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), strip_debug_ops=True) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['userDefinedMetadata']['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) weights_manifest = model_json['weightsManifest'] self.assertCountEqual(weights_manifest[0]['paths'], ['group1-shard1of1.bin']) self.assertIn('weights', weights_manifest[0]) self.assertTrue( glob.glob(os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_v1(self): self._create_saved_model_v1() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR)) weights = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{ 'dtype': 'float32', 'name': 'w', 'shape': [2, 2] }] }] tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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'] self.assertEqual(weights_manifest, weights) # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_skip_op_check(self): self._create_unsupported_saved_model() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), skip_op_check=True) weights = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{ 'dtype': 'float32', 'name': 'statefulpartitionedcall_args_1', 'shape': [2, 2] }, { 'dtype': 'float32', 'name': 'StatefulPartitionedCall/MatrixDiag', 'shape': [2, 2, 2] }] }] tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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'] self.assertEqual(weights_manifest, weights) self.assertTrue( glob.glob(os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_sharded(self): self._create_saved_model() model_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # Do initial conversion without sharding. tf_saved_model_conversion_v2.convert_tf_saved_model( model_path, tfjs_path) weight_files = glob.glob(os.path.join(tfjs_path, '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 2 shards after conversion. weight_shard_size_bytes = int(optimized_total_weight * 0.8) tfjs_path = os.path.join(self._tmp_dir, 'sharded_model') # Convert Saved Model again with shard argument set. tf_saved_model_conversion_v2.convert_tf_saved_model( model_path, tfjs_path, weight_shard_size_bytes=weight_shard_size_bytes) weight_files = sorted(glob.glob(os.path.join(tfjs_path, 'group*.bin'))) self.assertEqual(len(weight_files), 2) weight_file_sizes = [os.path.getsize(f) for f in weight_files] self.assertEqual(sum(weight_file_sizes), optimized_total_weight) self.assertLess(weight_file_sizes[1], weight_file_sizes[0])
def test_convert_saved_model_v1_with_hashtable(self): self._create_saved_model_v1_with_hashtable() input_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) output_dir = os.path.join(input_dir, 'js') tf_saved_model_conversion_v2.convert_tf_saved_model( input_dir, output_dir) expected_weights_manifest = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{ 'dtype': 'float32', 'name': 'w', 'shape': [2, 2] }, { 'dtype': 'string', 'name': 'Const', 'shape': [1] }, { 'dtype': 'int32', 'name': 'Const_1', 'shape': [1] }] }] tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'js') # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) self.assertTrue(model_json['modelInitializer']) for node in model_json['modelTopology']['node']: if node['name'] == 'ToFloat' and node['op'] == 'Placeholder': self.assertEqual( node['attr']['shape'], { 'shape': { 'dim': [{ 'size': '-1' }, { 'size': '2' }, { 'size': '2' }] } }) weights_manifest = model_json['weightsManifest'] self.assertEqual(weights_manifest, expected_weights_manifest) # 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 test_convert_saved_model(self): self._create_saved_model() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR)) weights = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{ 'dtype': 'float32', 'name': 'StatefulPartitionedCall/mul', 'shape': [] }] }] tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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'] self.assertEqual(len(weights_manifest), len(weights)) if sys.version_info[0] < 3: self.assertItemsEqual(weights_manifest[0]['paths'], weights[0]['paths']) self.assertItemsEqual(weights_manifest[0]['weights'], weights[0]['weights']) else: self.assertCountEqual(weights_manifest[0]['paths'], weights[0]['paths']) self.assertCountEqual(weights_manifest[0]['weights'], weights[0]['weights'])
def test_convert_saved_model_strip_debug_ops(self): self._create_saved_model_with_debug_ops() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), strip_debug_ops=True) weights = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{ 'dtype': 'float32', 'name': 'add', 'shape': [2, 2] }] }] tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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'] self.assertEqual(weights_manifest, weights) self.assertTrue( glob.glob(os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_with_unfusable_prelu(self): self._create_saved_model_with_unfusable_prelu() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR)) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['userDefinedMetadata']['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) nodes = model_json['modelTopology']['node'] prelu_op = None for node in nodes: if node['op'] == 'Prelu': prelu_op = node break self.assertTrue(prelu_op) # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_with_control_flow(self): self._create_saved_model_with_control_flow() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR)) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['userDefinedMetadata']['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) weights_manifest = model_json['weightsManifest'] self.assertCountEqual(weights_manifest[0]['paths'], ['group1-shard1of1.bin']) self.assertIn('weights', weights_manifest[0]) # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def export(path: str) -> None: torch_path = f"{path}.pth" onnx_path = f"{path}.onnx" tensorflow_path = f"{path}.pb" tensorflowjs_path = path model = LeNet5(10) model.load_state_dict(torch.load(torch_path)) model = model.eval().cpu() torch.onnx.export( model, torch.zeros((1, 1, 28, 28)).float(), onnx_path, input_names=["img"], output_names=["pred"], do_constant_folding=True, export_params=True, opset_version=10, verbose=True, ) model = onnx.load(onnx_path) prepare(model).export_graph(tensorflow_path) convert_tf_saved_model(tensorflow_path, tensorflowjs_path)
def dispatch_keras_h5_to_tfjs_graph_model_conversion( h5_path, output_dir=None, quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, use_structured_outputs_names=False, weight_shard_size_bytes=1024 * 1024 * 4, control_flow_v2=False, experiments=False, metadata=None): """ Convert a keras HDF5-format model to tfjs GraphModel artifacts. Args: h5_path: Path to the HDF5-format file that contains the model saved from keras or tf.keras. output_dir: The destination to which the tfjs GraphModel artifacts will be written. quantization_dtype_map: A mapping from dtype (`uint8`, `uint16`, `float16`) to weights. The weight mapping supports wildcard substitution. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to allow unsupported debug ops. use_structured_outputs_names: Bool whether output of graph model will follow the structured_outputs format. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. experiments: Bool enable experimental features. metadata: User defined metadata map. """ if not os.path.exists(h5_path): raise ValueError('Nonexistent path to HDF5 file: %s' % h5_path) if os.path.isdir(h5_path): raise ValueError( 'Expected path to point to an HDF5 file, but it points to a ' 'directory: %s' % h5_path) temp_savedmodel_dir = tempfile.mktemp(suffix='.savedmodel') model = tf.keras.models.load_model(h5_path, compile=False) model.save(temp_savedmodel_dir, include_optimizer=False, save_format='tf') # NOTE(cais): This cannot use `tf.compat.v1` because # `convert_tf_saved_model()` works only in v2. tf_saved_model_conversion_v2.convert_tf_saved_model( temp_savedmodel_dir, output_dir, signature_def='serving_default', saved_model_tags='serve', quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, use_structured_outputs_names=use_structured_outputs_names, weight_shard_size_bytes=weight_shard_size_bytes, control_flow_v2=control_flow_v2, experiments=experiments, metadata=metadata) # Clean up the temporary SavedModel directory. shutil.rmtree(temp_savedmodel_dir)
def save_tfjs_from_torch(artifacts_dir, example_input, model, model_file_name): model_file_path = join(artifacts_dir, model_file_name) os.makedirs(model_file_path, exist_ok=True) torch.onnx.export(model.cpu(), example_input, join(model_file_path, 'model.onnx'), export_params=True, opset_version=11) model_onnx = onnx.load(join(model_file_path, 'model.onnx')) model_tf = prepare(model_onnx) model_tf.export_graph(join(model_file_path, 'model')) tf_saved_model_conversion_v2.convert_tf_saved_model(join(model_file_path, 'model'), model_file_path, skip_op_check=True) rmtree(join(model_file_path, 'model')) os.remove(join(model_file_path, 'model.onnx'))
def test_convert_saved_model_with_frozen_file(self): self._create_saved_model() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), frozen_graph_dir=os.path.join(self._tmp_dir, SAVED_MODEL_DIR)) frozen_file_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'model.json.frozen') # Check model.json.frozen exist. self.assertTrue(glob.glob(frozen_file_path))
def test_convert_saved_model_structured_outputs_false(self): self._create_saved_model_with_structured_outputs() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR)) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # Check model.json and weights manifest. with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f: model_json = json.load(f) self.assertIs(model_json.get('userDefinedMetadata'), None)
def test_convert_saved_model_with_prelu(self): self._create_saved_model_with_prelu() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR) ) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['userDefinedMetadata']['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) nodes = model_json['modelTopology']['node'] prelu_op = None fused_op = None depthwise_fused_op = None for node in nodes: if node['op'] == 'Prelu': prelu_op = node if node['op'] == '_FusedConv2D': fused_op = node if node['op'] == graph_rewrite_util.FUSED_DEPTHWISE_CONV2D: depthwise_fused_op = node self.assertTrue(prelu_op is None) self.assertIsNot(fused_op, None) self.assertIsNot(depthwise_fused_op, None) fused_ops = list(map(base64.b64decode, fused_op['attr']['fused_ops']['list']['s'])) self.assertEqual(fused_ops, [b'BiasAdd', b'Prelu']) self.assertEqual(fused_op['attr']['num_args']['i'], '2') depthwise_fused_ops = list( map(base64.b64decode, depthwise_fused_op['attr']['fused_ops']['list']['s'])) self.assertEqual(depthwise_fused_ops, [b'BiasAdd', b'Prelu']) self.assertEqual(depthwise_fused_op['attr']['num_args']['i'], '2') # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def make_tfjs_export(tflite_saved_model_dir, export_dir): if os.path.exists(export_dir): log('TF-JS export already exists in {}, skipping TF-JS export'.format( export_dir)) return # Make a TF-JS model from the TF-Lite SavedModel export. log('Making TF-JS model ...') os.makedirs(export_dir) tfjs_saved_model_converter.convert_tf_saved_model(tflite_saved_model_dir, export_dir) log('Done')
def test_convert_saved_model_with_control_flow(self): self._create_saved_model_with_control_flow() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR) ) weights = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{'dtype': 'int32', 'shape': [], 'name': 'StatefulPartitionedCall/while/loop_counter'}, {'dtype': 'int32', 'shape': [], 'name': 'StatefulPartitionedCall/while/maximum_iterations' }, {'dtype': 'int32', 'shape': [], 'name': 'StatefulPartitionedCall/while/cond/_3/mod/y'}, {'dtype': 'int32', 'shape': [], 'name': 'StatefulPartitionedCall/while/cond/_3/Equal/y'}, {'dtype': 'int32', 'shape': [], 'name': 'StatefulPartitionedCall/while/body/_4/add_1/y'}, {'name': 'StatefulPartitionedCall/add/y', 'dtype': 'int32', 'shape': []}]}] tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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'] self.assertEqual(len(weights_manifest), len(weights)) if sys.version_info[0] < 3: self.assertItemsEqual(weights_manifest[0]['paths'], weights[0]['paths']) self.assertItemsEqual(weights_manifest[0]['weights'], weights[0]['weights']) else: self.assertCountEqual(weights_manifest[0]['paths'], weights[0]['paths']) self.assertCountEqual(weights_manifest[0]['weights'], weights[0]['weights']) # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def dispatch_keras_h5_to_tfjs_graph_model_conversion( h5_path, output_dir=None, quantization_dtype=None, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, control_flow_v2=False): """ Convert a keras HDF5-format model to tfjs GraphModel artifacts. Args: h5_path: Path to the HDF5-format file that contains the model saved from keras or tf.keras. output_dir: The destination to which the tfjs GraphModel artifacts will be written. quantization_dtype: The quantized data type to store the weights in (Default: `None`). skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to allow unsupported debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. """ if not os.path.exists(h5_path): raise ValueError('Nonexistent path to HDF5 file: %s' % h5_path) if os.path.isdir(h5_path): raise ValueError( 'Expected path to point to an HDF5 file, but it points to a ' 'directory: %s' % h5_path) temp_savedmodel_dir = tempfile.mktemp(suffix='.savedmodel') model = tf.keras.models.load_model(h5_path, compile=False) model.save(temp_savedmodel_dir, include_optimizer=False, save_format='tf') # NOTE(cais): This cannot use `tf.compat.v1` because # `convert_tf_saved_model()` works only in v2. tf_saved_model_conversion_v2.convert_tf_saved_model( temp_savedmodel_dir, output_dir, signature_def='serving_default', saved_model_tags='serve', quantization_dtype=quantization_dtype, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, control_flow_v2=control_flow_v2) # Clean up the temporary SavedModel directory. shutil.rmtree(temp_savedmodel_dir)
def test_convert_saved_model_with_control_flow_v2(self): self._create_saved_model_with_control_flow() tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) tf_saved_model_conversion_v2.convert_tf_saved_model( tfjs_path, tfjs_path, control_flow_v2=True ) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) weights_manifest = model_json['weightsManifest'] self.assertCountEqual(weights_manifest[0]['paths'], ['group1-shard1of1.bin']) self.assertIn('weights', weights_manifest[0]) add_y_weight = None for weight in weights_manifest[0]['weights']: if 'add/y' in weight['name']: add_y_weight = weight self.assertIsNot(add_y_weight, None) self.assertFalse(add_y_weight['name'].startswith('add/y')) nodes = model_json['modelTopology']['node'] while_op = None for node in nodes: self.assertNotIn('Merge', node['op']) self.assertNotIn('Switch', node['op']) if node['op'] == 'StatelessWhile': while_op = node self.assertIsNot(while_op, None) # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_with_metadata(self): self._create_saved_model() metadata_json = {'a': 1} tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), metadata={'key': metadata_json}) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # Check model.json and weights manifest. with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f: model_json = json.load(f) self.assertEqual(metadata_json, model_json['userDefinedMetadata']['key'])
def test_convert_saved_model_with_fused_depthwise_conv2d(self): self._create_saved_model_with_fusable_depthwise_conv2d() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR) ) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['userDefinedMetadata']['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) nodes = model_json['modelTopology']['node'] fusedOp = None for node in nodes: self.assertNotIn('BatchNorm', node['op']) self.assertNotIn('Relu', node['op']) self.assertNotIn('BiasAdd', node['op']) if node['op'] == graph_rewrite_util.FUSED_DEPTHWISE_CONV2D: fusedOp = node self.assertIsNot(fusedOp, None) self.assertIsNot(fusedOp['attr']['dilations'], None) self.assertIsNot(fusedOp['attr']['strides'], None) self.assertEqual( base64.b64decode(fusedOp['attr']['fused_ops']['list']['s'][0]), b'BiasAdd') self.assertEqual( base64.b64decode(fusedOp['attr']['fused_ops']['list']['s'][1]), b'Relu') # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model(self): self._create_saved_model() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR) ) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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'] self.assertCountEqual(weights_manifest[0]['paths'], ['group1-shard1of1.bin']) self.assertIn('weights', weights_manifest[0])
def dispatch_keras_h5_to_tfjs_graph_model_conversion(h5_path, output_dir=None, quantization_dtype=None, skip_op_check=False, strip_debug_ops=False): """ Convert a keras HDF5-format model to tfjs GraphModel artifacts. Args: h5_path: Path to the HDF5-format file that contains the model saved from keras or tf.keras. output_dir: The destination to which the tfjs GraphModel artifacts will be written. quantization_dtype: The quantized data type to store the weights in (Default: `None`). skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to allow unsupported debug ops. """ if not os.path.exists(h5_path): raise ValueError('Nonexistent path to HDF5 file: %s' % h5_path) if os.path.isdir(h5_path): raise ValueError( 'Expected path to point to an HDF5 file, but it points to a ' 'directory: %s' % h5_path) temp_savedmodel_dir = tempfile.mktemp(suffix='.savedmodel') model = keras.models.load_model(h5_path) keras.experimental.export_saved_model(model, temp_savedmodel_dir, serving_only=True) # NOTE(cais): This cannot use `tf.compat.v1` because # `convert_tf_saved_model()` works only in v2. tf_saved_model_conversion_v2.convert_tf_saved_model( temp_savedmodel_dir, output_dir, signature_def='serving_default', saved_model_tags='serve', quantization_dtype=quantization_dtype, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops) # Clean up the temporary SavedModel directory. shutil.rmtree(temp_savedmodel_dir)
def save_tfjs(model, combined_model_name): combined_model_name_dir = f'{tmpdir}/tfjs-models/{combined_model_name}' os.makedirs(combined_model_name_dir, exist_ok=True) example_input = torch.randn(1, 1, 176, requires_grad=False) torch.onnx.export(model.cpu(), example_input, f'{combined_model_name_dir}/model.onnx', export_params=True, opset_version=11) onnx_model = onnx.load(f'{combined_model_name_dir}/model.onnx') tf_model = prepare(onnx_model) tf_model.export_graph(f'{combined_model_name_dir}/model') tf_saved_model_conversion_v2.convert_tf_saved_model( f'{combined_model_name_dir}/model', combined_model_name_dir, skip_op_check=True) rmtree(f'{combined_model_name_dir}/model') os.remove(f'{combined_model_name_dir}/model.onnx')
def test_convert_saved_model_skip_op_check(self): self._create_unsupported_saved_model() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), skip_op_check=True ) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) weights_manifest = model_json['weightsManifest'] self.assertCountEqual(weights_manifest[0]['paths'], ['group1-shard1of1.bin']) self.assertIn('weights', weights_manifest[0]) self.assertTrue( glob.glob( os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_with_fused_conv2d(self): self._create_saved_model_with_fusable_conv2d() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR) ) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) nodes = model_json['modelTopology']['node'] fusedOp = None for node in nodes: self.assertTrue(not 'BatchNorm' in node['op']) self.assertTrue(not 'Relu' in node['op']) self.assertTrue(not 'BiasAdd' in node['op']) if node['op'] == '_FusedConv2D': fusedOp = node self.assertTrue(fusedOp is not None) self.assertEqual( base64.b64decode(fusedOp['attr']['fused_ops']['list']['s'][0]), b'BiasAdd') self.assertEqual( base64.b64decode(fusedOp['attr']['fused_ops']['list']['s'][1]), b'Relu') # 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(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*')))
def test_convert_saved_model_structured_outputs_true(self): self._create_saved_model_with_structured_outputs() tf_saved_model_conversion_v2.convert_tf_saved_model( os.path.join(self._tmp_dir, SAVED_MODEL_DIR), os.path.join(self._tmp_dir, SAVED_MODEL_DIR), use_structured_outputs_names=True) tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) # 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']) self.assertIsNot(model_json['modelTopology']['versions'], None) signature = model_json['signature'] self.assertIsNot(signature, None) self.assertIsNot(signature['inputs'], None) self.assertIsNot(signature['outputs'], None) self.assertEqual( ["a", "b", "c"], model_json['userDefinedMetadata']['structuredOutputKeys'])
def convert(arguments): args = get_arg_parser().parse_args(arguments) if args.show_version: print('\ntensorflowjs %s\n' % version.version) print('Dependency versions:') print(' keras %s' % keras.__version__) print(' tensorflow %s' % tf.__version__) return if not args.input_path: raise ValueError( 'Missing input_path argument. For usage, use the --help flag.') if not args.output_path: raise ValueError( 'Missing output_path argument. For usage, use the --help flag.') weight_shard_size_bytes = 1024 * 1024 * 4 if args.weight_shard_size_bytes: if args.output_format != common.TFJS_LAYERS_MODEL: raise ValueError( 'The --weight_shard_size_bytes flag is only supported under ' 'output_format=tfjs_layers_model.') weight_shard_size_bytes = args.weight_shard_size_bytes if args.input_path is None: raise ValueError('Error: The input_path argument must be set. ' 'Run with --help flag for usage information.') input_format, output_format = _standardize_input_output_formats( args.input_format, args.output_format) quantization_dtype = ( quantization.QUANTIZATION_BYTES_TO_DTYPES[args.quantization_bytes] if args.quantization_bytes else None) if (args.signature_name and input_format not in (common.TF_SAVED_MODEL, common.TF_HUB_MODEL)): raise ValueError( 'The --signature_name flag is applicable only to "tf_saved_model" and ' '"tf_hub" input format, but the current input format is ' '"%s".' % input_format) # TODO(cais, piyu): More conversion logics can be added as additional # branches below. if (input_format == common.KERAS_MODEL and output_format == common.TFJS_LAYERS_MODEL): dispatch_keras_h5_to_tfjs_layers_model_conversion( args.input_path, output_dir=args.output_path, quantization_dtype=quantization_dtype, split_weights_by_layer=args.split_weights_by_layer) elif (input_format == common.KERAS_MODEL and output_format == common.TFJS_GRAPH_MODEL): dispatch_keras_h5_to_tfjs_graph_model_conversion( args.input_path, output_dir=args.output_path, quantization_dtype=quantization_dtype, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) elif (input_format == common.KERAS_SAVED_MODEL and output_format == common.TFJS_LAYERS_MODEL): dispatch_keras_saved_model_to_tensorflowjs_conversion( args.input_path, args.output_path, quantization_dtype=quantization_dtype, split_weights_by_layer=args.split_weights_by_layer) elif (input_format == common.TF_SAVED_MODEL and output_format == common.TFJS_GRAPH_MODEL): tf_saved_model_conversion_v2.convert_tf_saved_model( args.input_path, args.output_path, signature_def=args.signature_name, saved_model_tags=args.saved_model_tags, quantization_dtype=quantization_dtype, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) elif (input_format == common.TF_HUB_MODEL and output_format == common.TFJS_GRAPH_MODEL): tf_saved_model_conversion_v2.convert_tf_hub_module( args.input_path, args.output_path, args.signature_name, args.saved_model_tags, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) elif (input_format == common.TFJS_LAYERS_MODEL and output_format == common.KERAS_MODEL): dispatch_tensorflowjs_to_keras_h5_conversion(args.input_path, args.output_path) elif (input_format == common.TFJS_LAYERS_MODEL and output_format == common.KERAS_SAVED_MODEL): dispatch_tensorflowjs_to_keras_saved_model_conversion( args.input_path, args.output_path) elif (input_format == common.TFJS_LAYERS_MODEL and output_format == common.TFJS_LAYERS_MODEL): dispatch_tensorflowjs_to_tensorflowjs_conversion( args.input_path, args.output_path, quantization_dtype=_parse_quantization_bytes( args.quantization_bytes), weight_shard_size_bytes=weight_shard_size_bytes) elif (input_format == common.TFJS_LAYERS_MODEL and output_format == common.TFJS_GRAPH_MODEL): dispatch_tfjs_layers_model_to_tfjs_graph_conversion( args.input_path, args.output_path, quantization_dtype=_parse_quantization_bytes( args.quantization_bytes), skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) else: raise ValueError( 'Unsupported input_format - output_format pair: %s - %s' % (input_format, output_format))