# -*- encoding:utf-8 -*- import logging import os import numpy as np import tensorflow as tf from google.protobuf import text_format from tensorflow.python.framework import ops from tensorflow.python.platform.gfile import GFile # from tensorflow.python.saved_model import constants from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model.loader_impl import SavedModelLoader from easy_rec.python.utils import conditional from easy_rec.python.utils import constant from easy_rec.python.utils import embedding_utils from easy_rec.python.utils import proto_util EMBEDDING_INITIALIZERS = 'embedding_initializers' class MetaGraphEditor: def __init__(self, lookup_lib_path, saved_model_dir, redis_url=None, redis_passwd=None, redis_timeout=0, redis_cache_names=[], oss_path=None, oss_endpoint=None, oss_ak=None, oss_sk=None, oss_timeout=0, meta_graph_def=None, norm_name_to_ids=None, incr_update_params=None, debug_dir=''): self._lookup_op = tf.load_op_library(lookup_lib_path) self._debug_dir = debug_dir self._verbose = debug_dir != '' if saved_model_dir: tags = ['serve'] loader = SavedModelLoader(saved_model_dir) saver, _ = loader.load_graph(tf.get_default_graph(), tags, None) meta_graph_def = loader.get_meta_graph_def_from_tags(tags) else: assert meta_graph_def, 'either saved_model_dir or meta_graph_def must be set' tf.reset_default_graph() from tensorflow.python.framework import meta_graph meta_graph.import_scoped_meta_graph_with_return_elements( meta_graph_def, clear_devices=True) # tf.train.import_meta_graph(meta_graph_def) self._meta_graph_version = meta_graph_def.meta_info_def.meta_graph_version self._signature_def = meta_graph_def.signature_def[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] if self._verbose: debug_out_path = os.path.join(self._debug_dir, 'meta_graph_raw.txt') with GFile(debug_out_path, 'w') as fout: fout.write(text_format.MessageToString(meta_graph_def, as_utf8=True)) self._meta_graph_def = meta_graph_def self._old_node_num = len(self._meta_graph_def.graph_def.node) self._all_graph_nodes = None self._all_graph_node_flags = None self._restore_tensor_node = None self._restore_shard_node = None self._restore_all_node = [] self._lookup_outs = None self._feature_names = None self._embed_names = None self._embed_name_to_ids = norm_name_to_ids self._is_cache_from_redis = [] self._redis_cache_names = redis_cache_names self._embed_ids = None self._embed_dims = None self._embed_sizes = None self._embed_combiners = None self._redis_url = redis_url self._redis_passwd = redis_passwd self._redis_timeout = redis_timeout self._oss_path = oss_path self._oss_endpoint = oss_endpoint self._oss_ak = oss_ak self._oss_sk = oss_sk self._oss_timeout = oss_timeout self._incr_update_params = incr_update_params # increment update placeholders self._embedding_update_inputs = {} self._embedding_update_outputs = {} self._dense_update_inputs = {} self._dense_update_outputs = {} @property def sparse_update_inputs(self): return self._embedding_update_inputs @property def sparse_update_outputs(self): return self._embedding_update_outputs @property def dense_update_inputs(self): return self._dense_update_inputs @property def dense_update_outputs(self): return self._dense_update_outputs @property def graph_def(self): return self._meta_graph_def.graph_def @property def signature_def(self): return self._signature_def @property def meta_graph_version(self): return self._meta_graph_version def init_graph_node_clear_flags(self): graph_def = self._meta_graph_def.graph_def self._all_graph_nodes = [n for n in graph_def.node] self._all_graph_node_flags = [True for n in graph_def.node] def _get_share_embed_name(self, x, embed_names): """Map share embedding tensor names to embed names. Args: x: string, embedding tensor names, such as: input_layer_1/shared_embed_1/field16_shared_embedding input_layer_1/shared_embed_2/field17_shared_embedding input_layer/shared_embed_wide/field15_shared_embedding input_layer/shared_embed_wide_1/field16_shared_embedding embed_names: all the optional embedding_names Return: one element in embed_names, such as: input_layer_1/shared_embed input_layer_1/shared_embed input_layer/shared_embed_wide input_layer/shared_embed_wide """ assert x.endswith('_shared_embedding') name_toks = x.split('/') name_toks = name_toks[:-1] tmp = name_toks[-1] tmp = tmp.split('_') try: int(tmp[-1]) name_toks[-1] = '_'.join(tmp[:-1]) except Exception: pass tmp_name = '/'.join(name_toks[1:]) sel_embed_name = '' for embed_name in embed_names: tmp_toks = embed_name.split('/') tmp_toks = tmp_toks[1:] embed_name_sub = '/'.join(tmp_toks) if tmp_name == embed_name_sub: assert not sel_embed_name, 'confusions encountered: %s %s' % ( x, ','.join(embed_names)) sel_embed_name = embed_name assert sel_embed_name, '%s not find in shared_embeddings: %s' % ( tmp_name, ','.join(embed_names)) return sel_embed_name def _find_embed_combiners(self, norm_embed_names): """Find embedding lookup combiner methods. Args: norm_embed_names: normalized embedding names Return: list: combiner methods for each features: sum, mean, sqrtn """ embed_combiners = {} embed_combine_node_cts = {} combiner_map = { 'SparseSegmentSum': 'sum', 'SparseSegmentMean': 'mean', 'SparseSegmentSqrtN': 'sqrtn' } for node in self._meta_graph_def.graph_def.node: if node.op in combiner_map: norm_name, _ = proto_util.get_norm_embed_name(node.name) embed_combiners[norm_name] = combiner_map[node.op] embed_combine_node_cts[norm_name] = embed_combine_node_cts.get( norm_name, 0) + 1 elif node.op == 'RealDiv' and len(node.input) == 2: # for tag feature with weights, and combiner == mean if 'SegmentSum' in node.input[0] and 'SegmentSum' in node.input[1]: norm_name, _ = proto_util.get_norm_embed_name(node.name) embed_combiners[norm_name] = 'mean' embed_combine_node_cts[norm_name] = embed_combine_node_cts.get( norm_name, 0) + 1 elif node.op == 'SegmentSum': norm_name, _ = proto_util.get_norm_embed_name(node.name) # avoid overwrite RealDiv results if norm_name not in embed_combiners: embed_combiners[norm_name] = 'sum' embed_combine_node_cts[norm_name] = embed_combine_node_cts.get( norm_name, 0) + 1 return [embed_combiners[x] for x in norm_embed_names] def _find_lookup_indices_values_shapes(self): # use the specific _embedding_weights/SparseReshape to find out # lookup inputs: indices, values, dense_shape, weights indices = {} values = {} shapes = {} def _get_output_shape(graph_def, input_name): out_id = 0 if ':' in input_name: node_name, out_id = input_name.split(':') out_id = int(out_id) else: node_name = input_name for node in graph_def.node: if node.name == node_name: return node.attr['_output_shapes'].list.shape[out_id] return None for node in self._meta_graph_def.graph_def.node: if '_embedding_weights/SparseReshape' in node.name: if node.op == 'SparseReshape': # embed_name, _ = proto_util.get_norm_embed_name(node.name, self._verbose) fea_name, _ = proto_util.get_norm_embed_name(node.name, self._verbose) for tmp_input in node.input: tmp_shape = _get_output_shape(self._meta_graph_def.graph_def, tmp_input) if '_embedding_weights/Cast' in tmp_input: continue elif len(tmp_shape.dim) == 2: indices[fea_name] = tmp_input elif len(tmp_shape.dim) == 1: shapes[fea_name] = tmp_input elif node.op == 'Identity': fea_name, _ = proto_util.get_norm_embed_name(node.name, self._verbose) values[fea_name] = node.input[0] return indices, values, shapes def _find_lookup_weights(self): weights = {} for node in self._meta_graph_def.graph_def.node: if '_weighted_by_' in node.name and 'GatherV2' in node.name: has_sparse_reshape = False for tmp_input in node.input: if 'SparseReshape' in tmp_input: has_sparse_reshape = True if has_sparse_reshape: continue if len(node.input) != 3: continue # try to find nodes with weights # input_layer/xxx_weighted_by_yyy_embedding/xxx_weighted_by_yyy_embedding_weights/GatherV2_[0-9] # which has three inputs: # input_layer/xxx_weighted_by_yyy_embedding/xxx_weighted_by_yyy_embedding_weights/Reshape_1 # DeserializeSparse_1 (this is the weight) # input_layer/xxx_weighted_by_yyy_embedding/xxx_weighted_by_yyy_embedding_weights/GatherV2_4/axis fea_name, _ = proto_util.get_norm_embed_name(node.name, self._verbose) for tmp_input in node.input: if '_weighted_by_' not in tmp_input: weights[fea_name] = tmp_input return weights def _find_embed_names_and_dims(self, norm_embed_names): # get embedding dimensions from Variables embed_dims = {} embed_sizes = {} embed_is_kv = {} for node in self._meta_graph_def.graph_def.node: if 'embedding_weights' in node.name and node.op in [ 'VariableV2', 'KvVarHandleOp' ]: tmp = node.attr['shape'].shape.dim[-1].size tmp2 = 1 for x in node.attr['shape'].shape.dim[:-1]: tmp2 = tmp2 * x.size embed_name, _ = proto_util.get_norm_embed_name(node.name, self._verbose) assert embed_name is not None,\ 'fail to get_norm_embed_name(%s)' % node.name embed_dims[embed_name] = tmp embed_sizes[embed_name] = tmp2 embed_is_kv[embed_name] = 1 if node.op == 'KvVarHandleOp' else 0 # get all embedding dimensions, note that some embeddings # are shared by multiple inputs, so the names should be # transformed all_embed_dims = [] all_embed_names = [] all_embed_sizes = [] all_embed_is_kv = [] for x in norm_embed_names: if x in embed_dims: all_embed_names.append(x) all_embed_dims.append(embed_dims[x]) all_embed_sizes.append(embed_sizes[x]) all_embed_is_kv.append(embed_is_kv[x]) elif x.endswith('_shared_embedding'): tmp_embed_name = self._get_share_embed_name(x, embed_dims.keys()) all_embed_names.append(tmp_embed_name) all_embed_dims.append(embed_dims[tmp_embed_name]) all_embed_sizes.append(embed_sizes[tmp_embed_name]) all_embed_is_kv.append(embed_is_kv[tmp_embed_name]) return all_embed_names, all_embed_dims, all_embed_sizes, all_embed_is_kv def find_lookup_inputs(self): logging.info('Extract embedding_lookup inputs') indices, values, shapes = self._find_lookup_indices_values_shapes() weights = self._find_lookup_weights() for fea in shapes.keys(): logging.info('Lookup Input[%s]: indices=%s values=%s shapes=%s' % (fea, indices[fea], values[fea], shapes[fea])) graph = tf.get_default_graph() def _get_tensor_by_name(tensor_name): if ':' not in tensor_name: tensor_name = tensor_name + ':0' return graph.get_tensor_by_name(tensor_name) lookup_input_values = [] lookup_input_indices = [] lookup_input_shapes = [] lookup_input_weights = [] for key in values.keys(): tmp_val, tmp_ind, tmp_shape = values[key], indices[key], shapes[key] lookup_input_values.append(_get_tensor_by_name(tmp_val)) lookup_input_indices.append(_get_tensor_by_name(tmp_ind)) lookup_input_shapes.append(_get_tensor_by_name(tmp_shape)) if key in weights: tmp_w = weights[key] lookup_input_weights.append(_get_tensor_by_name(tmp_w)) else: lookup_input_weights.append([]) # get embedding combiners self._embed_combiners = self._find_embed_combiners(values.keys()) # get embedding dimensions self._embed_names, self._embed_dims, self._embed_sizes, self._embed_is_kv\ = self._find_embed_names_and_dims(values.keys()) if not self._embed_name_to_ids: embed_name_uniq = list(set(self._embed_names)) self._embed_name_to_ids = { t: tid for tid, t in enumerate(embed_name_uniq) } self._embed_ids = [ int(self._embed_name_to_ids[x]) for x in self._embed_names ] self._is_cache_from_redis = [ proto_util.is_cache_from_redis(x, self._redis_cache_names) for x in self._embed_names ] # normalized feature names self._feature_names = list(values.keys()) return lookup_input_indices, lookup_input_values, lookup_input_shapes,\ lookup_input_weights def add_lookup_op(self, lookup_input_indices, lookup_input_values, lookup_input_shapes, lookup_input_weights): logging.info('add custom lookup operation to lookup embeddings from redis') self._lookup_outs = [None for i in range(len(lookup_input_values))] for i in range(len(lookup_input_values)): if lookup_input_values[i].dtype == tf.int32: lookup_input_values[i] = tf.to_int64(lookup_input_values[i]) for i in range(len(self._lookup_outs)): i_1 = i + 1 self._lookup_outs[i] = self._lookup_op.kv_lookup( lookup_input_indices[i:i_1], lookup_input_values[i:i_1], lookup_input_shapes[i:i_1], lookup_input_weights[i:i_1], url=self._redis_url, password=self._redis_passwd, timeout=self._redis_timeout, combiners=self._embed_combiners[i:i_1], embedding_dims=self._embed_dims[i:i_1], embedding_names=self._embed_ids[i:i_1], cache=self._is_cache_from_redis, version=self._meta_graph_version)[0] meta_graph_def = tf.train.export_meta_graph() if self._verbose: debug_path = os.path.join(self._debug_dir, 'graph_raw.txt') with GFile(debug_path, 'w') as fout: fout.write( text_format.MessageToString( self._meta_graph_def.graph_def, as_utf8=True)) return meta_graph_def def add_oss_lookup_op(self, lookup_input_indices, lookup_input_values, lookup_input_shapes, lookup_input_weights): logging.info('add custom lookup operation to lookup embeddings from oss') place_on_cpu = os.getenv('place_embedding_on_cpu') place_on_cpu = eval(place_on_cpu) if place_on_cpu else False with conditional(place_on_cpu, ops.device('/CPU:0')): for i in range(len(lookup_input_values)): if lookup_input_values[i].dtype == tf.int32: lookup_input_values[i] = tf.to_int64(lookup_input_values[i]) # N = len(lookup_input_indices) # self._lookup_outs = [ None for _ in range(N) ] # for i in range(N): # i_1 = i + 1 # self._lookup_outs[i] = self._lookup_op.oss_read_kv( # lookup_input_indices[i:i_1], # lookup_input_values[i:i_1], # lookup_input_shapes[i:i_1], # lookup_input_weights[i:i_1], # osspath=self._oss_path, # endpoint=self._oss_endpoint, # ak=self._oss_ak, # sk=self._oss_sk, # timeout=self._oss_timeout, # combiners=self._embed_combiners[i:i_1], # embedding_dims=self._embed_dims[i:i_1], # embedding_ids=self._embed_ids[i:i_1], # embedding_is_kv=self._embed_is_kv[i:i_1], # shared_name='embedding_lookup_res', # name='embedding_lookup_fused/lookup')[0] self._lookup_outs = self._lookup_op.oss_read_kv( lookup_input_indices, lookup_input_values, lookup_input_shapes, lookup_input_weights, osspath=self._oss_path, endpoint=self._oss_endpoint, ak=self._oss_ak, sk=self._oss_sk, timeout=self._oss_timeout, combiners=self._embed_combiners, embedding_dims=self._embed_dims, embedding_ids=self._embed_ids, embedding_is_kv=self._embed_is_kv, shared_name='embedding_lookup_res', name='embedding_lookup_fused/lookup') N = np.max([int(x) for x in self._embed_ids]) + 1 uniq_embed_ids = [x for x in range(N)] uniq_embed_dims = [0 for x in range(N)] uniq_embed_combiners = ['mean' for x in range(N)] uniq_embed_is_kvs = [0 for x in range(N)] for embed_id, embed_combiner, embed_is_kv, embed_dim in zip( self._embed_ids, self._embed_combiners, self._embed_is_kv, self._embed_dims): uniq_embed_combiners[embed_id] = embed_combiner uniq_embed_is_kvs[embed_id] = embed_is_kv uniq_embed_dims[embed_id] = embed_dim lookup_init_op = self._lookup_op.oss_init( osspath=self._oss_path, endpoint=self._oss_endpoint, ak=self._oss_ak, sk=self._oss_sk, combiners=uniq_embed_combiners, embedding_dims=uniq_embed_dims, embedding_ids=uniq_embed_ids, embedding_is_kv=uniq_embed_is_kvs, N=N, shared_name='embedding_lookup_res', name='embedding_lookup_fused/init') ops.add_to_collection(EMBEDDING_INITIALIZERS, lookup_init_op) if self._incr_update_params is not None: # all sparse variables are updated by a single custom operation message_ph = tf.placeholder(tf.int8, [None], name='incr_update/message') embedding_update = self._lookup_op.embedding_update( message=message_ph, shared_name='embedding_lookup_res', name='embedding_lookup_fused/embedding_update') self._embedding_update_inputs['incr_update/sparse/message'] = message_ph self._embedding_update_outputs[ 'incr_update/sparse/embedding_update'] = embedding_update # dense variables are updated one by one dense_name_to_ids = embedding_utils.get_dense_name_to_ids() for x in ops.get_collection(constant.DENSE_UPDATE_VARIABLES): dense_var_id = dense_name_to_ids[x.op.name] dense_input_name = 'incr_update/dense/%d/input' % dense_var_id dense_output_name = 'incr_update/dense/%d/output' % dense_var_id dense_update_input = tf.placeholder( tf.float32, x.get_shape(), name=dense_input_name) self._dense_update_inputs[dense_input_name] = dense_update_input dense_assign_op = tf.assign(x, dense_update_input) self._dense_update_outputs[dense_output_name] = dense_assign_op meta_graph_def = tf.train.export_meta_graph() if self._verbose: debug_path = os.path.join(self._debug_dir, 'graph_raw.txt') with GFile(debug_path, 'w') as fout: fout.write( text_format.MessageToString( self._meta_graph_def.graph_def, as_utf8=True)) return meta_graph_def def bytes2str(self, x): if bytes == str: return x else: try: return x.decode('utf-8') except Exception: # in case of some special chars in protobuf return str(x) def clear_meta_graph_embeding(self, meta_graph_def): logging.info('clear meta graph embedding_weights') def _clear_embedding_in_meta_collect(meta_graph_def, collect_name): tmp_vals = [ x for x in meta_graph_def.collection_def[collect_name].bytes_list.value if 'embedding_weights' not in self.bytes2str(x) ] meta_graph_def.collection_def[collect_name].bytes_list.ClearField('value') for tmp_v in tmp_vals: meta_graph_def.collection_def[collect_name].bytes_list.value.append( tmp_v) _clear_embedding_in_meta_collect(meta_graph_def, 'model_variables') _clear_embedding_in_meta_collect(meta_graph_def, 'trainable_variables') _clear_embedding_in_meta_collect(meta_graph_def, 'variables') # clear Kv(pai embedding variable) ops in meta_info_def.stripped_op_list.op kept_ops = [ x for x in meta_graph_def.meta_info_def.stripped_op_list.op if x.name not in [ 'InitializeKvVariableOp', 'KvResourceGather', 'KvResourceImportV2', 'KvVarHandleOp', 'KvVarIsInitializedOp', 'ReadKvVariableOp' ] ] meta_graph_def.meta_info_def.stripped_op_list.ClearField('op') meta_graph_def.meta_info_def.stripped_op_list.op.extend(kept_ops) for tmp_op in meta_graph_def.meta_info_def.stripped_op_list.op: if tmp_op.name == 'SaveV2': for tmp_id, tmp_attr in enumerate(tmp_op.attr): if tmp_attr.name == 'has_ev': tmp_op.attr.remove(tmp_attr) break def clear_meta_collect(self, meta_graph_def): drop_meta_collects = [] for key in meta_graph_def.collection_def: val = meta_graph_def.collection_def[key] if val.HasField('node_list'): if 'embedding_weights' in val.node_list.value[ 0] and 'easy_rec' not in val.node_list.value[0]: drop_meta_collects.append(key) elif key == 'saved_model_assets': drop_meta_collects.append(key) for key in drop_meta_collects: meta_graph_def.collection_def.pop(key) def remove_embedding_weights_and_update_lookup_outputs(self): def _should_drop(name): if '_embedding_weights' in name: if self._verbose: logging.info('[SHOULD_DROP] %s' % name) return True logging.info('remove embedding_weights node in graph_def.node') logging.info( 'and replace the old embedding_lookup outputs with new lookup_op outputs' ) for tid, node in enumerate(self._all_graph_nodes): # drop the nodes if _should_drop(node.name): self._all_graph_node_flags[tid] = False else: for i in range(len(node.input)): if _should_drop(node.input[i]): input_name, _ = proto_util.get_norm_embed_name( node.input[i], self._verbose) print('REPLACE:' + node.input[i] + '=>' + input_name) input_name = self._lookup_outs[self._feature_names.index( input_name)].name if input_name.endswith(':0'): input_name = input_name.replace(':0', '') node.input[i] = input_name # drop by ids def _drop_by_ids(self, tmp_obj, key, drop_ids): keep_vals = [ x for i, x in enumerate(getattr(tmp_obj, key)) if i not in drop_ids ] tmp_obj.ClearField(key) getattr(tmp_obj, key).extend(keep_vals) def clear_save_restore(self): """Clear save restore ops. save/restore_all need save/restore_shard as input save/restore_shard needs save/Assign_[0-N] as input save/Assign_[0-N] needs save/RestoreV2 as input save/RestoreV2 use save/RestoreV2/tensor_names and save/RestoreV2/shape_and_slices as input edit [ save/RestoreV2/tensor_names save/RestoreV2/shape_and_slices save/RestoreV2 save/restore_shard ] """ for tid, node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if node.name == 'save/RestoreV2/tensor_names': self._restore_tensor_node = node break # assert self._restore_tensor_node is not None, 'save/RestoreV2/tensor_names is not found' if self._restore_tensor_node: drop_ids = [] for tmp_id, tmp_name in enumerate( self._restore_tensor_node.attr['value'].tensor.string_val): if 'embedding_weights' in self.bytes2str(tmp_name): drop_ids.append(tmp_id) self._drop_by_ids(self._restore_tensor_node.attr['value'].tensor, 'string_val', drop_ids) keep_node_num = len( self._restore_tensor_node.attr['value'].tensor.string_val) logging.info( 'update self._restore_tensor_node: string_val keep_num = %d drop_num = %d' % (keep_node_num, len(drop_ids))) self._restore_tensor_node.attr['value'].tensor.tensor_shape.dim[ 0].size = keep_node_num self._restore_tensor_node.attr['_output_shapes'].list.shape[0].dim[ 0].size = keep_node_num logging.info( 'update save/RestoreV2, drop tensor_shapes, _output_shapes, related to embedding_weights' ) self._restore_shard_node = None for node_id, node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if node.name == 'save/RestoreV2/shape_and_slices': node.attr['value'].tensor.tensor_shape.dim[0].size = keep_node_num node.attr['_output_shapes'].list.shape[0].dim[0].size = keep_node_num self._drop_by_ids(node.attr['value'].tensor, 'string_val', drop_ids) elif node.name == 'save/RestoreV2': self._drop_by_ids(node.attr['_output_shapes'].list, 'shape', drop_ids) self._drop_by_ids(node.attr['dtypes'].list, 'type', drop_ids) elif node.name == 'save/restore_shard': self._restore_shard_node = node elif node.name.startswith('save/restore_all'): self._restore_all_node.append(node) def clear_save_assign(self): logging.info( 'update save/Assign, drop tensor_shapes, _output_shapes, related to embedding_weights' ) # edit save/Assign drop_save_assigns = [] all_kv_drop = [] for tid, node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if node.op == 'Assign' and 'save/Assign' in node.name and \ 'embedding_weights' in node.input[0]: drop_save_assigns.append('^' + node.name) self._all_graph_node_flags[tid] = False elif 'embedding_weights/ConcatPartitions/concat' in node.name: self._all_graph_node_flags[tid] = False elif node.name.endswith('/embedding_weights') and node.op == 'Identity': self._all_graph_node_flags[tid] = False elif 'save/KvResourceImportV2' in node.name and node.op == 'KvResourceImportV2': drop_save_assigns.append('^' + node.name) self._all_graph_node_flags[tid] = False elif 'KvResourceImportV2' in node.name: self._all_graph_node_flags[tid] = False elif 'save/Const' in node.name and node.op == 'Const': if '_class' in node.attr and len(node.attr['_class'].list.s) > 0: const_name = node.attr['_class'].list.s[0] if not isinstance(const_name, str): const_name = const_name.decode('utf-8') if 'embedding_weights' in const_name: self._all_graph_node_flags[tid] = False elif 'ReadKvVariableOp' in node.name and node.op == 'ReadKvVariableOp': all_kv_drop.append(node.name) self._all_graph_node_flags[tid] = False elif node.op == 'Assign' and 'save/Assign' in node.name: # update node(save/Assign_[0-N])'s input[1] by the position of # node.input[0] in save/RestoreV2/tensor_names # the outputs of save/RestoreV2 is connected to save/Assign tmp_id = [ self.bytes2str(x) for x in self._restore_tensor_node.attr['value'].tensor.string_val ].index(node.input[0]) if tmp_id != 0: tmp_input2 = 'save/RestoreV2:%d' % tmp_id else: tmp_input2 = 'save/RestoreV2' if tmp_input2 != node.input[1]: if self._verbose: logging.info("update save/Assign[%s]'s input from %s to %s" % (node.name, node.input[1], tmp_input2)) node.input[1] = tmp_input2 # save/restore_all need save/restore_shard as input # save/restore_shard needs save/Assign_[0-N] as input # save/Assign_[0-N] needs save/RestoreV2 as input if self._restore_shard_node: for tmp_input in drop_save_assigns: self._restore_shard_node.input.remove(tmp_input) if self._verbose: logging.info('drop restore_shard input: %s' % tmp_input) elif len(self._restore_all_node) > 0: for tmp_input in drop_save_assigns: for tmp_node in self._restore_all_node: if tmp_input in tmp_node.input: tmp_node.input.remove(tmp_input) if self._verbose: logging.info('drop %s input: %s' % (tmp_node.name, tmp_input)) break def clear_save_v2(self): """Clear SaveV2 ops. save/Identity need [ save/MergeV2Checkpoints, save/control_dependency ] as input. Save/MergeV2Checkpoints need [save/MergeV2Checkpoints/checkpoint_prefixes] as input. Save/MergeV2Checkpoints/checkpoint_prefixes need [ save/ShardedFilename, save/control_dependency ] as input. save/control_dependency need save/SaveV2 as input. save/SaveV2 input: [ save/SaveV2/tensor_names, save/SaveV2/shape_and_slices ] edit save/SaveV2 save/SaveV2/shape_and_slices save/SaveV2/tensor_names. """ logging.info('update save/SaveV2 input shape, _output_shapes, tensor_shape') save_drop_ids = [] for tid, node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if node.name == 'save/SaveV2' and node.op == 'SaveV2': for tmp_id, tmp_input in enumerate(node.input): if '/embedding_weights' in tmp_input: save_drop_ids.append(tmp_id) diff_num = len(node.input) - len(node.attr['dtypes'].list.type) self._drop_by_ids(node, 'input', save_drop_ids) save_drop_ids = [x - diff_num for x in save_drop_ids] self._drop_by_ids(node.attr['dtypes'].list, 'type', save_drop_ids) if 'has_ev' in node.attr: del node.attr['has_ev'] for node in self._all_graph_nodes: if node.name == 'save/SaveV2/shape_and_slices' and node.op == 'Const': # _output_shapes # size # string_val node.attr['_output_shapes'].list.shape[0].dim[0].size -= len( save_drop_ids) node.attr['value'].tensor.tensor_shape.dim[0].size -= len(save_drop_ids) self._drop_by_ids(node.attr['value'].tensor, 'string_val', save_drop_ids) elif node.name == 'save/SaveV2/tensor_names': # tensor_names may not have the same order as save/SaveV2/shape_and_slices tmp_drop_ids = [ tmp_id for tmp_id, tmp_val in enumerate( node.attr['value'].tensor.string_val) if 'embedding_weights' in self.bytes2str(tmp_val) ] # attr['value'].tensor.string_val # tensor_shape # size assert len(save_drop_ids) == len(save_drop_ids) node.attr['_output_shapes'].list.shape[0].dim[0].size -= len( tmp_drop_ids) node.attr['value'].tensor.tensor_shape.dim[0].size -= len(tmp_drop_ids) self._drop_by_ids(node.attr['value'].tensor, 'string_val', tmp_drop_ids) def clear_initialize(self): """Clear initialization ops. */read(Identity) depend on [*(VariableV2)] */Assign depend on [*/Initializer/*, *(VariableV2)] drop embedding_weights initialization nodes */embedding_weights/part_x [,/Assign,/read] */embedding_weights/part_1/Initializer/truncated_normal [,/shape,/mean,/stddev,/TruncatedNormal,/mul] """ logging.info('Remove Initialization nodes for embedding_weights') for tid, node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if 'embedding_weights' in node.name and 'Initializer' in node.name: self._all_graph_node_flags[tid] = False elif 'embedding_weights' in node.name and 'Assign' in node.name: self._all_graph_node_flags[tid] = False elif 'embedding_weights' in node.name and node.op == 'VariableV2': self._all_graph_node_flags[tid] = False elif 'embedding_weights' in node.name and node.name.endswith( '/read') and node.op == 'Identity': self._all_graph_node_flags[tid] = False elif 'embedding_weights' in node.name and node.op == 'Identity': node_toks = node.name.split('/') node_tok = node_toks[-1] if 'embedding_weights_' in node_tok: node_tok = node_tok[len('embedding_weights_'):] try: int(node_tok) self._all_graph_node_flags[tid] = False except Exception: pass def clear_embedding_variable(self): # for pai embedding variable, we drop some special nodes for tid, node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if node.op in [ 'ReadKvVariableOp', 'KvVarIsInitializedOp', 'KvVarHandleOp' ]: self._all_graph_node_flags[tid] = False # there maybe some nodes depend on the dropped nodes, they are dropped as well def drop_dependent_nodes(self): drop_names = [ tmp_node.name for tid, tmp_node in enumerate(self._all_graph_nodes) if not self._all_graph_node_flags[tid] ] while True: more_drop_names = [] for tid, tmp_node in enumerate(self._all_graph_nodes): if not self._all_graph_node_flags[tid]: continue if len(tmp_node.input) > 0 and tmp_node.input[0] in drop_names: logging.info('drop dependent node: %s depend on %s' % (tmp_node.name, tmp_node.input[0])) self._all_graph_node_flags[tid] = False more_drop_names.append(tmp_node.name) drop_names = more_drop_names if not drop_names: break def edit_graph(self): # the main entrance lookup_input_indices, lookup_input_values, lookup_input_shapes,\ lookup_input_weights = self.find_lookup_inputs() # add lookup op to the graph self._meta_graph_def = self.add_lookup_op(lookup_input_indices, lookup_input_values, lookup_input_shapes, lookup_input_weights) self.clear_meta_graph_embeding(self._meta_graph_def) self.clear_meta_collect(self._meta_graph_def) self.init_graph_node_clear_flags() self.remove_embedding_weights_and_update_lookup_outputs() # save/RestoreV2 self.clear_save_restore() # save/Assign self.clear_save_assign() # save/SaveV2 self.clear_save_v2() self.clear_initialize() self.clear_embedding_variable() self.drop_dependent_nodes() self._meta_graph_def.graph_def.ClearField('node') self._meta_graph_def.graph_def.node.extend([ x for tid, x in enumerate(self._all_graph_nodes) if self._all_graph_node_flags[tid] ]) logging.info('old node number = %d' % self._old_node_num) logging.info('node number = %d' % len(self._meta_graph_def.graph_def.node)) if self._verbose: debug_dump_path = os.path.join(self._debug_dir, 'graph.txt') with GFile(debug_dump_path, 'w') as fout: fout.write(text_format.MessageToString(self.graph_def, as_utf8=True)) debug_dump_path = os.path.join(self._debug_dir, 'meta_graph.txt') with GFile(debug_dump_path, 'w') as fout: fout.write( text_format.MessageToString(self._meta_graph_def, as_utf8=True)) def edit_graph_for_oss(self): # the main entrance lookup_input_indices, lookup_input_values, lookup_input_shapes,\ lookup_input_weights = self.find_lookup_inputs() # add lookup op to the graph self._meta_graph_def = self.add_oss_lookup_op(lookup_input_indices, lookup_input_values, lookup_input_shapes, lookup_input_weights) self.clear_meta_graph_embeding(self._meta_graph_def) self.clear_meta_collect(self._meta_graph_def) self.init_graph_node_clear_flags() self.remove_embedding_weights_and_update_lookup_outputs() # save/RestoreV2 self.clear_save_restore() # save/Assign self.clear_save_assign() # save/SaveV2 self.clear_save_v2() self.clear_initialize() self.clear_embedding_variable() self.drop_dependent_nodes() self._meta_graph_def.graph_def.ClearField('node') self._meta_graph_def.graph_def.node.extend([ x for tid, x in enumerate(self._all_graph_nodes) if self._all_graph_node_flags[tid] ]) logging.info('old node number = %d' % self._old_node_num) logging.info('node number = %d' % len(self._meta_graph_def.graph_def.node)) if self._verbose: debug_dump_path = os.path.join(self._debug_dir, 'graph.txt') with GFile(debug_dump_path, 'w') as fout: fout.write(text_format.MessageToString(self.graph_def, as_utf8=True)) debug_dump_path = os.path.join(self._debug_dir, 'meta_graph.txt') with GFile(debug_dump_path, 'w') as fout: fout.write( text_format.MessageToString(self._meta_graph_def, as_utf8=True))