博客

详解TensorFlow的 tf.nn.embedding_lookup_sparse 函数：查找稀疏张量中的嵌入矩阵

  • Post category:Python

tf.nn.embedding_lookup_sparse 函数是 TensorFlow 中用于实现稀疏张量乘法的函数，可以根据稀疏张量中的索引在词向量矩阵中查找对应的词向量，其具体作用是通过查找词向量表并获得稀疏张量对应的向量，以获得向量表示的映射。以下是该函数的函数定义：

tf.nn.embedding_lookup_sparse(
    params,
    sp_ids,
    sp_weights,
    combiner='mean',
    partition_strategy='mod',
    name=None,
    max_norm=None
)

该函数的参数含义如下：

  • params: 词向量矩阵，是一个形状为 (vocab_size, embedding_size) 的矩阵。
  • sp_ids: 一个 SparseTensor 对象，包含待查找的稀疏索引，该张量必须是一个稀疏矩阵，其 tensorShape 可以与 params 的前i维形状为相同，第i维必定为1，sparse_values为各个sparseIndex对应的embbedding id。
  • sp_weights: 一个 SparseTensor 对象，包含待查找的稀疏索引的权重。如果是二元分类问题，则为实际标签的路径数值；如果是多元分类问题，则为第i条训练样本，第j个预测标签存在的概率分布值。
  • combiner: 稀疏张量与参数矩阵相乘后的合并方法，可选项为 mean, sqrtn, and sum，默认为 mean。其中，mean 将在稀疏索引的权重下计算每个向量的平均值，sqrtn 对每个向量进行归一化，以确保向量的长度平均值为1，sum 将对应的向量相加。
  • partition_strategy: 选择如何将操作分解到不同设备上。默认为 mod，即使用下标模运算符进行分区，还可选择 min 和 div。
  • name: 操作的名称。
  • max_norm: 对词向量进行规范化的最大范数，以使其不超过给定的值。

下面提供两个示例，以说明该函数的用法：

示例一

在模型中使用该函数，将分割后的特征行（稀疏矩阵）桶化为 n 个 dense feature。下面是使用示例：

# 定义一个具有 5 个 dense feature 的线性模型
def linear_model(features, labels):
  vocab_size = 1000  # 词表大小
  embedding_size = 10  # 词向量维度
  n = 5  # 分段数

  # 定义词向量矩阵
  embedding = tf.Variable(tf.truncated_normal([vocab_size, embedding_size]))

  # 对稀疏张量进行分段和桶化
  buckets = tf.feature_column.bucketized_column(features['sparse_feature'], boundaries=[...])
  embeddings = tf.feature_column.embedding_column(buckets, dimension=embedding_size)
  sparse_feature_matrix = tf.feature_column.input_layer(features, [embeddings])

  # 获取稀疏矩阵对应的词向量
  sparse_feature = tf.nn.embedding_lookup_sparse(
      embedding, sp_ids=sparse_feature_matrix, sp_weights=None, combiner='mean'
  )

  # 将稀疏向量与其他 dense feature 进行连接
  dense_features = tf.concat([sparse_feature, features['other_features']], axis=1)

  # 构建线性模型
  logits = tf.layers.dense(dense_features, 1, name='linear_model')
  loss = tf.losses.sigmoid_cross_entropy(labels, logits)

  return loss

... # 定义数据输入等其它步骤

# 训练模型
estimator = tf.estimator.Estimator(linear_model)
estimator.train(input_fn=input_fn, steps=10000)

在上述代码中，我们定义了一个有 5 个 dense feature 的线性模型，其中的稀疏矩阵使用了 tf.nn.embedding_lookup_sparse。首先，我们定义了事先定义好的词向量矩阵，然后对特征进行分段和桶化，得到一个稀疏矩阵 sparse_feature_matrix，该矩阵由使用桶化后的索引为行，对象的词向量编号为列，每个元素表示该特定索引在该对象里出现的次数组成。在调用 tf.nn.embedding_lookup_sparse 函数时，指定 sp_ids 为稀疏矩阵，sp_weights 为 None，表示不考虑权重。combiner 参数设置为 mean，即在稀疏索引权重的基础上，对查找到的向量进行平均计算，得到与该矩阵对应的词向量矩阵。最后，对 sparse_feature 与其他 dense feature 进行连接即可得到最终的输入。

示例二

使用词向量计算获取两个字符串的相似度。具体来说，对于每个字符串，利用词向量的平均值获取其表示向量，然后计算这两个向量之间的余弦相似度。使用示例如下：

import tensorflow_hub as hub

embed = hub.Module("https://tfhub.dev/google/nnlm-en-dim128/1")

# 获取两个句子的词向量表示
message_embeddings = embed(["A long sentence.", "single-word"])

# 计算两个句子之间的相似度
similarity = tf.nn.softmax(tf.matmul(message_embeddings, tf.transpose(message_embeddings)))

在上面的代码中，我们使用 TensorFlow Hub 提供的预先训练好的词向量模型 nnlm-en-dim128。该模型能够将每个词映射到一个 128 维的向量空间，因此，对于一个长句子，我们可以使用 tf.reduce_mean() 函数，沿着指定轴对词向量进行平均，获得该句子的向量表示。然后，对于两个句子，我们使用 tf.nn.embedding_lookup_sparse 函数，获取其向量表示，然后计算两个向量之间的余弦相似度，从而确定相似度。最后使用softmax处理得到两句话的相似度。

Tags: