TensorFlow、Dataset API 和 flat_map 操作

Question

我在使用tf.contrib.data.Dataset API 时遇到了困难，想知道你们中的一些人是否可以提供帮助。我想把word2vec的整个skip-gram预处理变成这个范式来玩API一点点，它涉及到以下操作：

1. 令牌序列是动态加载的（避免一次将所有数据集加载到内存中），比如说我们以Stream开始（理解为Scala的方式，所有数据不在内存中，而是在访问时加载）需要）的令牌序列：seq_tokens。
2. 从这些seq_tokens 中的任何一个中，我们使用一个返回元组列表(token, context) 的python 函数提取skip-grams。
3. 选择tokens 列作为特征，contexts 列作为标签。

在伪代码中使其更清晰，它看起来像上面。我们应该利用框架并行系统而不是自己加载数据，所以我会做一些事情，比如首先在内存中加载序列的索引，然后加载序列（在 map 内，因此如果不是所有行都是同步处理，数据是异步加载的，没有 OOM 可担心），并在这些令牌序列上应用一个函数，这将创建需要展平的不同数量的 skip-gram。为此，我将正式以data 的形式结束(#lines=number of skip-grams generated, #columns=2)。

data = range(1:N)
  .map(i => load(i): Seq[String]) // load: Int -> Seq[String] loads dynamically a sequence of tokens (sequences have varying length)
  .flat_map(s => skip_gram(s)) // skip_gram: Seq[String] -> Seq[(String, String)] with output length
features = data[0] // features
lables = data[1] // labels

我曾天真地尝试使用Dataset 的 API 这样做，但我卡住了，我可以这样做：

iterator = (
  tf.contrib.data.Dataset.range(N)
  .map(lambda i: tf.py_func(load_data, [i], [tf.int32, tf.int32])) // (1)
  .flat_map(?) // (2)
  .make_one_shot_iterator()
)

(1) TensorFlow 对此并不满意，因为加载的序列长度不同...

(2) 还没有完成skip-gram部分...我实际上只是想调用一个python函数来计算skip-gram的序列（可变大小）并将其展平，以便如果返回类型是一个矩阵，那么每一行应该理解为输出Dataset的换行。

如果有人有任何想法，非常感谢，如果我忘记提及有用的精度，请不要犹豫......

Answer 1

我只是在实现同样的事情；这是我解决它的方法：

dataset = tf.data.TextLineDataset(filename)

if mode == ModeKeys.TRAIN:
  dataset = dataset.shuffle(buffer_size=batch_size * 100)

dataset = dataset.flat_map(lambda line: string_to_skip_gram(line))

dataset = dataset.batch(batch_size)

在我的数据集中，我将每一行视为独立的，因此我不必担心跨越多行的上下文。

因此，我通过函数string_to_skip_gram 对每一行进行平面映射，该函数返回一个长度取决于行中标记数的Dataset。

string_to_skip_gram 使用tf.py_func 将行转换为一系列标记，由 ID 表示（使用方法tokenize_str）：

def string_to_skip_gram(line):
  def handle_line(line):
    token_ids = tokenize_str(line)

    (features, labels) = skip_gram(token_ids)

    return np.array([features, labels], dtype=np.int64)

  res = tf.py_func(handle_line, [line], tf.int64)

  features = res[0]
  labels = res[1]

  return tf.data.Dataset.from_tensor_slices((features, labels))

最后，skip_gram 返回所有可能的上下文词和目标词的列表：

def skip_gram(token_ids):
  skip_window = 1

  features = []
  labels = []

  context_range = [i for i in range(-skip_window, skip_window + 1) if i != 0]

  for word_index in range(skip_window, len(token_ids) - skip_window):
    for context_word_offset in context_range:
      features.append(token_ids[word_index])
      labels.append(token_ids[word_index + context_word_offset])

  return features, labels

请注意，我在这里没有对上下文词进行抽样；只是使用它们。