(model_path, save_path, opt_config=None, data_type=BF16, calib_file=None)
| 556 | |
| 557 | |
| 558 | def optimize(model_path, save_path, opt_config=None, data_type=BF16, calib_file=None): |
| 559 | saved_model = loader_impl._parse_saved_model(model_path) |
| 560 | tags = saved_model.meta_graphs[0].meta_info_def.tags |
| 561 | with tf.Session() as sess: |
| 562 | meta_graph_def = tf.saved_model.loader.load(sess, tags, model_path) |
| 563 | signature_keys = list(meta_graph_def.signature_def.keys()) |
| 564 | signature_def = meta_graph_def.signature_def[signature_keys[0]] |
| 565 | model_inputs = [_nd(v.name) for v in signature_def.inputs.values()] |
| 566 | model_outputs = [_nd(v.name) for v in signature_def.outputs.values()] |
| 567 | init_op = loader_impl.get_init_op(meta_graph_def) |
| 568 | if init_op is not None: |
| 569 | model_outputs.append(init_op.name) |
| 570 | frozen_gdef = tf.graph_util.convert_variables_to_constants( |
| 571 | sess, sess.graph_def, model_outputs |
| 572 | ) |
| 573 | |
| 574 | # Embedding & Dense optimization |
| 575 | dense_opt_dict = dense_opt(sess, frozen_gdef, opt_config, data_type, calib_file) |
| 576 | variable_path = f'{model_path}/variables/variables' |
| 577 | embed_opt_dict = embedding_opt( |
| 578 | sess, frozen_gdef, opt_config, data_type, variable_path |
| 579 | ) |
| 580 | ev_dict = update_embedding_vars(sess) |
| 581 | if len(ev_dict) > 0: |
| 582 | global_step = tf.train.get_global_step() |
| 583 | model_outputs.append(_nd(global_step.name)) |
| 584 | if isinstance(global_step, tf.Variable): |
| 585 | sess.run(tf.variables_initializer([global_step])) |
| 586 | |
| 587 | def _extract_sub_graph(outputs): |
| 588 | graph_def = sess.graph.as_graph_def(add_shapes=True) |
| 589 | util.remove_underscore_class_attr(graph_def) |
| 590 | return tf.graph_util.extract_sub_graph(graph_def, outputs) |
| 591 | |
| 592 | def _save(save_path): |
| 593 | sub_graph_def = _extract_sub_graph(model_inputs + model_outputs) |
| 594 | node_names = [node.name for node in sub_graph_def.node] |
| 595 | variables = [v for v in get_all_variables() if _nd(v.name) in node_names] |
| 596 | init_name = tf.variables_initializer(variables).name |
| 597 | saver = tf.train.Saver(variables, sharded=True, allow_empty=True) |
| 598 | saver.save(sess, save_path, write_meta_graph=False, write_state=False) |
| 599 | return saver, init_name |
| 600 | |
| 601 | # Create Saver |
| 602 | tmp_path = tempfile.mkdtemp(dir='.') |
| 603 | variable_path = f'{tmp_path}/variables' |
| 604 | saver, init_name = _save(variable_path) |
| 605 | # Optimize embedding variables |
| 606 | ev_opt_dict, opt_variable_path = embedding_var_opt( |
| 607 | sess, frozen_gdef, opt_config, data_type, variable_path |
| 608 | ) |
| 609 | if len(ev_opt_dict) > 0: |
| 610 | saver, init_name = _save(variable_path) |
| 611 | variable_path = opt_variable_path |
| 612 | |
| 613 | saver_nodes = [ |
| 614 | saver.saver_def.restore_op_name, |
| 615 | _nd(saver.saver_def.filename_tensor_name), |
no test coverage detected