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

我正在使用编写输入管道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)

这个答案是为了回应塞尔索·弗兰卡 https://stackoverflow.com/users/8414280/celso-fran%C3%A7a的评论。

我确实找到了一种方法，但没有返回字典而是使用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