如何将 py_func 与返回 dict 的函数一起使用

Posted 2023-02-23

【中文标题】如何将 py_func 与返回 dict 的函数一起使用

【英文标题】：How to use py_func with a function that returns dict

【发布时间】：2018-08-05 19:31:05

【问题描述】：

我正在使用tf.data.Dataset 编写输入管道。我想使用 python 代码来加载和转换我的样本，代码返回一个张量字典。不幸的是，我看不到如何将其定义为传递给tf.py_func 的输出类型。

我有一个解决方法，我的函数返回张量列表而不是字典，但它使我的代码可读性降低，因为我在该字典中有 4 个键。

代码如下所示

file_list = ....

def load(file_name):
    return "image": np.zeros(...,dtype=np.float32),
           "label": 1.0 # there is more labels, in the original code

ds = tf.data.Dataset.from_tensor_slices(file_list)
ds.shuffle(...)
out_type = ['image':tf.float32, "label":tf.float32 ] # ???? 
ds.map(lambda x: tf.py_func(load, [x], out_type))

ds.batch(...)
ds.prefetch(1)

【问题讨论】：

tf.py_func 不支持返回字典的函数。所有 TensorFlow 操作都使用单个张量或它们的列表/元组作为输入和输出，我认为您没有办法将其作为列表返回。

以上评论并不完全正确，tf.data.Dataset 可以很好地与字典配合使用。

我也有同样的问题。你找到解决办法了吗？

如果您在使用 tf.data.Dataset 编写输入管道时只考虑使用 tf.py_func，则有一种解决方法。本质上，通过两个步骤：首先使用tf.py_func 映射到张量，然后使用tf.dataset.map 创建字典。阅读this question 了解具体示例：

@CelsoFrança 我发布了一个答案。完整代码见github.com/jonrosner/bert/blob/master/praktikum_nlp.py

【参考方案1】：

这个回答是为了回应Celso Franca的评论。

我确实找到了一种方法，但没有返回字典，而是使用tf_example.SerializeToString()。

这两个函数用于动态处理 BERT 输入。它工作得更好，为我节省了许多前期预处理时间，同时在训练过程中没有损失任何性能。

def _convert(label, text):
    """Decodes a csv-line to a TensorFlow Example, serialized as a string."""
    np_label = label.numpy()
    np_text = text.numpy()
    tokens_a = tokenizer.tokenize(np_text)
    # Account for [CLS] and [SEP] with "- 2"
    if len(tokens_a) > seq_length - 2:
        tokens_a = tokens_a[0: (seq_length - 2)]
    tokens = []
    segment_ids = []
    tokens.append("[CLS]")
    segment_ids.append(0)
    for token in tokens_a:
        tokens.append(token)
        segment_ids.append(0)
    tokens.append("[SEP]")
    segment_ids.append(0)

    input_ids = tokenizer.convert_tokens_to_ids(tokens)
    # The mask has 1 for real tokens and 0 for padding tokens. Only real
    # tokens are attended to.
    input_mask = [1] * len(input_ids)

    # Zero-pad up to the sequence length.
    while len(input_ids) < seq_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)

    assert len(input_ids) == seq_length
    assert len(input_mask) == seq_length
    assert len(segment_ids) == seq_length

    label_id = label_map[np_label]
    features = collections.OrderedDict()
    features["input_ids"] = create_int_feature(input_ids)
    features["input_mask"] = create_int_feature(input_mask)
    features["segment_ids"] = create_int_feature(segment_ids)
    features["label_ids"] = create_int_feature([label_id])
    features["is_real_example"] = create_int_feature([int(True)])
    tf_example = tf.train.Example(features=tf.train.Features(feature=features))
    # tf.py_function only accepts true tf datatypes like string
    return tf_example.SerializeToString()

  def _decode_record(record):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t
    return example

  def input_fn(params):
    """The actual input function."""
    filenames = tf.data.Dataset.list_files(file_pattern)
    label_col = processor.get_label_col()
    text_col = processor.get_text_col()
    d = filenames.apply(
      tf.contrib.data.parallel_interleave(
          lambda filename: tf.data.experimental.CsvDataset(filename,
            [tf.float32, tf.string],
            select_cols=[label_col, text_col],
            field_delim=delimiter,
            header=True),
          cycle_length=2))
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)
    d = d.map(lambda label, text: tf.py_function(_convert, [label, text], tf.string))
    d = d.map(_decode_record)
    d = d.batch(batch_size=params["batch_size"], drop_remainder=drop_remainder)
    return d

【讨论】：

我仍然使用tf.py_func，但我只在从map函数返回输入时将样本转换为字典。