分布式 tensorflow 源码解读 3：lookup.index_table_from_tensor

import tensorflow as tf
sess = tf.Session()vocabulary_list = tf.constant(["emerson", "lake", "palmer"])table = tf.contrib.lookup.index_table_from_tensor(        vocabulary_list=vocabulary_list, num_oov_buckets=10, default_value=-1)features = tf.constant(["emerson", "lake", "and", "palmer", "dad", "mom", "hello"])table.init.run(session=sess)ids = table.lookup(features)print(sess.run(ids))
返回值：[ 0  1 10  2  9  3  9]

def index_table_from_tensor(vocabulary_list,                            num_oov_buckets=0,                            default_value=-1,                            hasher_spec=FastHashSpec,                            dtype=dtypes.string,                            name=None):  """Returns a lookup table that converts a string tensor into int64 IDs.  This operation constructs a lookup table to convert tensor of strings into  int64 IDs. The mapping can be initialized from a string `vocabulary_list` 1-D  tensor where each element is a key and corresponding index within the tensor  is the value.
  Args:    vocabulary_list: A 1-D `Tensor` that specifies the mapping of keys to      indices. The type of this object must be castable to `dtype`.    num_oov_buckets: The number of out-of-vocabulary buckets.    default_value: The value to use for out-of-vocabulary feature values.      Defaults to -1.    hasher_spec: A `HasherSpec` to specify the hash function to use for      assignment of out-of-vocabulary buckets.    dtype: The type of values passed to `lookup`. Only string and integers are      supported.    name: A name for this op (optional).  Returns:    The lookup table to map an input `Tensor` to index `int64` `Tensor`.  Raises:    ValueError: If `vocabulary_list` is invalid.    ValueError: If `num_oov_buckets` is negative.  """
  “”“  类型检查代码  ”“”
  with ops.name_scope(name, "string_to_index") as feat_to_id_scope:        # 获取由传入的vocabulary_list构造的hash表，keys是vocabulary_list里面的元素，values是    # 【0， len(num_elements）-1】    keys = ops.convert_to_tensor(vocabulary_list)    num_elements = array_ops.size(keys)    values = math_ops.to_int64(math_ops.range(num_elements))
    shared_name = ""    with ops.name_scope(None, "hash_table") as hash_table_scope:      table_keys = math_ops.to_int64(keys) if keys.dtype.is_integer else keys      init = KeyValueTensorInitializer(          table_keys,          values,          table_keys.dtype.base_dtype,          dtypes.int64,          name="table_init")      table = HashTable(          init, default_value, shared_name=shared_name, name=hash_table_scope)    if num_oov_buckets:      table = IdTableWithHashBuckets(          table,          num_oov_buckets=num_oov_buckets,          hasher_spec=hasher_spec,          name=feat_to_id_scope,          key_dtype=dtype)    return table

class KeyValueTensorInitializer(TableInitializerBase):  """Table initializers given `keys` and `values` tensors."""
  def __init__(self, keys, values, key_dtype=None, value_dtype=None, name=None):    with ops.name_scope(name, "key_value_init", [keys, values]) as scope:      self._keys = ops.convert_to_tensor(keys, dtype=key_dtype, name="keys")      self._values = ops.convert_to_tensor(          values, dtype=value_dtype, name="values")      self._name = scope
    super(KeyValueTensorInitializer, self).__init__(self._keys.dtype,                                                    self._values.dtype)
  def initialize(self, table):    """Initializes the given `table` with `keys` and `values` tensors.    Args:      table: The table to initialize.    Returns:      The operation that initializes the table.    Raises:      TypeError: when the keys and values data types do not match the table      key and value data types.    """    _check_table_dtypes(table, self._keys.dtype, self._values.dtype)    with ops.name_scope(        self._name, values=(table.table_ref, self._keys,                            self._values)) as scope:        init_op = gen_lookup_ops.lookup_table_import_v2(            table.table_ref, self._keys, self._values, name=scope)        ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op)    return init_op可看出table的初始化操作是通过lookup_table_import_v2这个op来完成，下面就分析下这个op的相关c代码：
REGISTER_OP("LookupTableImportV2")    .Input("table_handle: resource")    .Input("keys: Tin")    .Input("values: Tout")    .Attr("Tin: type")    .Attr("Tout: type")    .SetShapeFn([](InferenceContext* c) {      ShapeHandle handle;      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
      ShapeHandle keys;      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &keys));      TF_RETURN_IF_ERROR(c->Merge(keys, c->input(2), &keys));      return Status::OK();    });
// Clear the table and insert data.class LookupTableImportOp : public OpKernel { public:  explicit LookupTableImportOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
  void Compute(OpKernelContext* ctx) override {    lookup::LookupInterface* table;    OP_REQUIRES_OK(ctx, GetLookupTable("table_handle", ctx, &table));        “”“    ....    ”“”
    const Tensor& keys = ctx->input(1);    const Tensor& values = ctx->input(2);    OP_REQUIRES_OK(ctx, table->CheckKeyAndValueTensorsForImport(keys, values));
    int memory_used_before = 0;    if (ctx->track_allocations()) {      memory_used_before = table->MemoryUsed();    }    OP_REQUIRES_OK(ctx, table->ImportValues(ctx, keys, values));    if (ctx->track_allocations()) {      ctx->record_persistent_memory_allocation(table->MemoryUsed() -                                               memory_used_before);    }  }};

class HashTable(InitializableLookupTableBase):
  def __init__(self, initializer, default_value, shared_name=None, name=None):    """Creates a non-initialized `HashTable` object.    Creates a table, the type of its keys and values are specified by the    initializer.    Before using the table you will have to initialize it. After initialization    the table will be immutable.    Args:      initializer: The table initializer to use. See `HashTable` kernel for        supported key and value types.      default_value: The value to use if a key is missing in the table.      shared_name: If non-empty, this table will be shared under        the given name across multiple sessions.      name: A name for the operation (optional).    Returns:      A `HashTable` object.    """    with ops.name_scope(name, "hash_table", (initializer,                                             default_value)) as scope:      table_ref = gen_lookup_ops.hash_table_v2(          shared_name=shared_name,          key_dtype=initializer.key_dtype,          value_dtype=initializer.value_dtype,          name=scope)
      super(HashTable, self).__init__(table_ref, default_value, initializer)      self._value_shape = self._default_value.get_shape()

class InitializableLookupTableBase(LookupInterface):  """Initializable lookup table interface.  An initializable lookup tables persist across different steps.  """
  def __init__(self, table_ref, default_value, initializer):    """Construct a table object from a table reference.    If requires a table initializer object (subclass of `TableInitializerBase`).    It provides the table key and value types, as well as the op to initialize    the table. The caller is responsible to execute the initialization op.    Args:      table_ref: The table reference, i.e. the output of the lookup table ops.      default_value: The value to use if a key is missing in the table.      initializer: The table initializer to use.    """        name = table_ref.op.name.split("/")[-1]    super(InitializableLookupTableBase,          self).__init__(initializer.key_dtype, initializer.value_dtype,                         name)    self._table_ref = table_ref    self._default_value = ops.convert_to_tensor(        default_value, dtype=self._value_dtype)    self._default_value.get_shape().merge_with(tensor_shape.scalar())    self._init = initializer.initialize(self)
  @property  def table_ref(self):    """Get the underlying table reference."""    return self._table_ref
  @property  def default_value(self):    """The default value of the table."""    return self._default_value
  @property  def init(self):    """The table initialization op."""    return self._init
  def lookup(self, keys, name=None):    """Looks up `keys` in a table, outputs the corresponding values.    The `default_value` is used for keys not present in the table.    """    key_tensor = keys    with ops.name_scope(name, "%s_Lookup" % self._name,                        (self._table_ref, key_tensor,                         self._default_value)) as scope:      values = gen_lookup_ops.lookup_table_find_v2(          self._table_ref, key_tensor, self._default_value, name=scope)
    values.set_shape(key_tensor.get_shape())    if isinstance(keys, sparse_tensor.SparseTensor):      return sparse_tensor.SparseTensor(keys.indices, values, keys.dense_shape)    else:      return values

lookup_table_find_v2
// Table lookup op. Perform the lookup operation on the given table.class LookupTableFindOp : public OpKernel { public:  explicit LookupTableFindOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
  void Compute(OpKernelContext* ctx) override {    lookup::LookupInterface* table;    OP_REQUIRES_OK(ctx, GetLookupTable("table_handle", ctx, &table));        “”“    ......    ”“”
    const Tensor& key = ctx->input(1);    const Tensor& default_value = ctx->input(2);    OP_REQUIRES_OK(ctx, table->CheckFindArguments(key, default_value));
    TensorShape output_shape = key.shape();    output_shape.RemoveLastDims(table->key_shape().dims());    output_shape.AppendShape(table->value_shape());    Tensor* out;    OP_REQUIRES_OK(ctx, ctx->allocate_output("values", output_shape, &out));
    OP_REQUIRES_OK(ctx, table->Find(ctx, key, out, default_value));  }};
REGISTER_KERNEL_BUILDER(Name("LookupTableFind").Device(DEVICE_CPU),                        LookupTableFindOp);REGISTER_KERNEL_BUILDER(Name("LookupTableFindV2").Device(DEVICE_CPU),                        LookupTableFindOp);
# table->Find调用的是HashTable类的DoFind方法Status DoFind(const Tensor& key, Tensor* value,                const Tensor& default_value) override {    const V default_val = default_value.flat<V>()(0);    const auto key_values = key.flat<K>();    auto value_values = value->flat<V>();
    for (int64 i = 0; i < key_values.size(); ++i) {      value_values(i) = gtl::FindWithDefault(          *table_, SubtleMustCopyIfIntegral(key_values(i)), default_val);    }    return Status::OK();  }

import tensorflow as tfnum_oov_buckets = 3input_tensor = tf.constant(["emerson", "lake", "palmer", "king", "crimnson"])table = tf.IdTableWithHashBuckets(      tf.HashTable(tf.TextFileIdTableInitializer(filename), default_value),      num_oov_buckets)out = table.lookup(input_tensor).table.init.run()print(out.eval())
#filename依次为"emerson", "lake", "palmer"

class IdTableWithHashBuckets(LookupInterface):  """String to Id table wrapper that assigns out-of-vocabulary keys to buckets.  For example, if an instance of `IdTableWithHashBuckets` is initialized with a  string-to-id table that maps:  * `emerson -> 0`  * `lake -> 1`  * `palmer -> 2`  The `IdTableWithHashBuckets` object will performs the following mapping:  * `emerson -> 0`  * `lake -> 1`  * `palmer -> 2`  * `<other term> -> bucket_id`, where bucket_id will be between `3` and  `3 + num_oov_buckets - 1`, calculated by:  `hash(<term>) % num_oov_buckets + vocab_size`  If input_tensor is `["emerson", "lake", "palmer", "king", "crimson"]`,  the lookup result is `[0, 1, 2, 4, 7]`.  If `table` is None, only out-of-vocabulary buckets are used.  """
  def __init__(self,               table,               num_oov_buckets,               hasher_spec=FastHashSpec,               name=None,               key_dtype=None):    """Construct a `IdTableWithHashBuckets` object.    Args:      table: Table that maps `tf.string` or `tf.int64` keys to `tf.int64` ids.      num_oov_buckets: Number of buckets to use for out-of-vocabulary keys.      hasher_spec: A `HasherSpec` to specify the hash function to use for        assignation of out-of-vocabulary buckets  (optional).      name: A name for the operation (optional).      key_dtype: Data type of keys passed to `lookup`. Defaults to        `table.key_dtype` if `table` is specified, otherwise `tf.string`.        Must be string or integer, and must be castable to `table.key_dtype`.    Raises:      ValueError: when `table` in None and `num_oov_buckets` is not positive.      TypeError: when `hasher_spec` is invalid.    """    # If a name ends with a '/' it is a "name scope", remove all trailing '/'    # characters to use as table name.        “”“    类型检查    “”“
    self._num_oov_buckets = num_oov_buckets    super(IdTableWithHashBuckets, self).__init__(key_dtype, dtypes.int64,                                                 name.split("/")[-1])
  @property  def init(self):    """The table initialization op."""    if self._table:      return self._table.init    with ops.name_scope(None, "init"):      return control_flow_ops.no_op()
  @property  def table_ref(self):    if self._table is not None:      return self._table.table_ref    return None
  def _get_string_to_hash_bucket_fn(self, hasher_spec):    """Returns the string_to_hash_bucket op to use based on `hasher_spec`."""    if not isinstance(hasher_spec, HasherSpec):      raise TypeError("hasher_spec must be of type HasherSpec %s" % hasher_spec)    if hasher_spec.hasher == "fasthash":      return string_ops.string_to_hash_bucket_fast    if hasher_spec.hasher == "legacy":      return string_ops.string_to_hash_bucket    if hasher_spec.hasher == "stronghash":      return functools.partial(          string_ops.string_to_hash_bucket_strong, key=hasher_spec.key)    raise ValueError("Unknown hasher %s" % hasher_spec.hasher)
  def lookup(self, keys, name=None):
    values = keys    if self._num_oov_buckets == 0:      ids = self._table.lookup(values, name=name)    else:      # TODO(yleon): Consider moving this functionality to its own kernel.      with ops.name_scope(name, "%s_Lookup" % self.name) as scope:        str_to_hash_bucket = self._get_string_to_hash_bucket_fn(            self._hasher_spec)        buckets = str_to_hash_bucket(            _as_string(values),            num_buckets=self._num_oov_buckets,            name="hash_bucket")        if self._table:          ids = self._table.lookup(values)          buckets = math_ops.add(buckets, self._table.size())          is_id_non_default = math_ops.not_equal(ids, self._table.default_value)          ids = array_ops.where(is_id_non_default, ids, buckets, name=scope)        else:          ids = buckets    return ids

 def lookup_v2(self, keys, name=None):
    values = keys       if self._num_oov_buckets == 0:       ids = self._table.lookup(values, name=name)    else:       with ops.name_scope(name, "%s_Lookup" % self.name) as scope:          str_to_hash_bucket = self._get_string_to_hash_bucket_fn(                    self._hasher_spec)          if self._table:              init_ids = self._table.lookup(values)                is_id_default_idx = array_ops.where(math_ops.equal(init_ids, self._table.default_value))                   hash_values = array_ops.gather(values, array_ops.squeeze(is_id_default_idx))               default_buckets = str_to_hash_bucket(                        _as_string(hash_values),                        num_buckets=self._num_oov_buckets,                        name="hash_bucket")              default_buckets = math_ops.add(default_buckets, self._table.size())              default_buckets = control_flow_ops.cond(gen_math_ops.equal(array_ops.size(default_buckets), 1),                                                            lambda: array_ops.expand_dims(default_buckets, axis=0),                                                            lambda: default_buckets)              ids = gen_array_ops.tensor_scatter_update(init_ids, is_id_default_idx, default_buckets)          else:              ids = str_to_hash_bucket(_as_string(values), num_buckets=self._num_oov_buckets, name="hash_bucket")    return ids