easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/utils/meta_graph_editor.py

@@ -0,0 +1,941 @@
+# -*- 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))