(session, graph_def, opt_config, data_type, variable_path)
| 367 | |
| 368 | |
| 369 | def embedding_opt(session, graph_def, opt_config, data_type, variable_path): |
| 370 | simple_graph = SimpleGraph(graph_def) |
| 371 | |
| 372 | def _get_gather_pattern(): |
| 373 | pl = list() |
| 374 | pl.append(SimpleNode('gather', 'GatherV2', ['read', '0', '1'], ['0'])) |
| 375 | pl.append(SimpleNode('read', 'Identity', ['embed'], ['gather'])) |
| 376 | pl.append(SimpleNode('embed', 'Const', [], ['read'])) |
| 377 | pattern_nodes = {node.name: node for node in pl} |
| 378 | return pattern_nodes, pl[0].name |
| 379 | |
| 380 | update_dict = dict() |
| 381 | pattern, first_key = _get_gather_pattern() |
| 382 | ptm_list = util.get_matched_pattern(simple_graph, pattern, first_key) |
| 383 | if opt_config: |
| 384 | ptm_list = [ptm for ptm in ptm_list if ptm['embed'] in opt_config] |
| 385 | for ptm in ptm_list: |
| 386 | embed_node = util.get_node_by_name(graph_def, simple_graph, ptm['embed']) |
| 387 | opt_dtype = opt_config.get(embed_node.name) if opt_config else data_type |
| 388 | if embed_node.name not in update_dict: |
| 389 | print(f'Optimize embedding to {opt_dtype}: {embed_node.name}') |
| 390 | # Add variables |
| 391 | fp32_data = util.get_const_value_by_name( |
| 392 | graph_def, ptm['embed'], simple_graph |
| 393 | ) |
| 394 | if opt_dtype == INT8: |
| 395 | int8_name = f'{embed_node.name}/int8_data' |
| 396 | int8_var = tf.get_variable(int8_name, fp32_data.shape, tf.int8) |
| 397 | scale_name = f'{embed_node.name}/int8_scale' |
| 398 | scale_var = tf.get_variable(scale_name, int8_var.shape[-1:], tf.float32) |
| 399 | update_dict[embed_node.name] = [int8_var, scale_var, opt_dtype] |
| 400 | elif opt_dtype in [BF16, FP16]: |
| 401 | tf_dtype = tf.bfloat16 if opt_dtype == BF16 else tf.float16 |
| 402 | f16_name = f'{embed_node.name}/{opt_dtype.lower()}_data' |
| 403 | f16_var = tf.get_variable(f16_name, fp32_data.shape, tf_dtype) |
| 404 | update_dict[embed_node.name] = [f16_var, opt_dtype] |
| 405 | else: |
| 406 | raise Exception(f'Unsupported data type: {opt_dtype}') |
| 407 | # Update Graph |
| 408 | gather_op = session.graph.get_operation_by_name(ptm['gather']) |
| 409 | opt_gather = tf.gather( |
| 410 | params=update_dict[embed_node.name][0], |
| 411 | indices=gather_op.inputs[1], |
| 412 | axis=gather_op.inputs[2], |
| 413 | batch_dims=gather_op.get_attr('batch_dims'), |
| 414 | name=f'{ptm["gather"]}/{opt_dtype.lower()}', |
| 415 | ) |
| 416 | cast_name = f'{ptm["gather"]}/cast_to_fp32' |
| 417 | update_tensor = tf.cast(opt_gather, dtype=tf.float32, name=cast_name) |
| 418 | if opt_dtype == INT8: |
| 419 | rescale_name = f'{ptm["gather"]}/rescale' |
| 420 | scale_var = update_dict[embed_node.name][1] |
| 421 | update_tensor = tf.multiply(update_tensor, scale_var, name=rescale_name) |
| 422 | update_op_inputs(session.graph, {ptm['gather']: update_tensor}) |
| 423 | |
| 424 | # Convert checkpoint |
| 425 | for opt_dtype in [INT8, BF16, FP16]: |
| 426 | opt_dict = {k: v for k, v in update_dict.items() if v[-1] == opt_dtype} |
no test coverage detected