Как вернуть словарь тензоров из функции tf.py_function? - PullRequest
Новая
       23

Как вернуть словарь тензоров из функции tf.py_function?

0 голосов
/ 04 апреля 2020

Обычно токенайзер трансформаторов кодирует входные данные в виде словаря. 

{"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32}

И для лучшей обработки производительности архивирования с большим набором данных рекомендуется использовать конвейер, включающий использование Dataset.map, чтобы применить функцию токенизатора к каждому элементу входного набора данных. Точно так же, как в учебном пособии по Tensorflow: Загрузить текст .

Однако tf.py_function (используется для переноса карты python функция) не поддерживает возврат словаря тензоров, как показано выше.

Например, , если токенизатор (кодировщик) в тексте Load возвращает следующий словарь: 

{
    "input_ids": [ 101, 13366,  2131,  1035,  6819,  2094,  1035,  102 ],
    "attention_mask": [ 1, 1, 1, 1, 1, 1, 1, 1 ]
}

как можно установить параметр Tout для tf.py_function чтобы получить нужный словарь тензоров: 

{
    'input_ids': <tf.Tensor: shape=(16,), dtype=int32, numpy = array(
    [ 101, 13366,  2131,  1035,  6819,  2094,  1035,  102 ], dtype=int32)>

    'attention_mask': <tf.Tensor: shape=(16,), dtype=int32, numpy=array(
     [ 1, 1, 1, 1, 1, 1, 1, 1 ], dtype=int32)>
}

?

1 Ответ

1 голос
/ 04 апреля 2020

tf.py_function не допускает python dict в качестве типа возврата. https://github.com/tensorflow/tensorflow/issues/36276

В качестве обходного пути в вашем случае вы можете выполнить преобразование данных в вашем py_function и затем вызвать другой файл tf.map без использования py_function для возврата словаря. 

def gen():
  yield 1

def process_data(x):
  return ([ 101, 13366,  2131,  1035,  6819,  2094,  1035,  102 ],
          [ 1, 1, 1, 1, 1, 1, 1, 1 ])

def create_dict(input_ids, attention_mask):
  return {"input_ids": tf.convert_to_tensor(input_ids),
          "attention_mask": tf.convert_to_tensor(attention_mask)}

ds = (tf.data.Dataset
      .from_generator(gen, (tf.int32))
      .map(lambda x: tf.py_function(process_data, inp=[x], 
                                    Tout=(tf.int32, tf.int32)))
      .map(create_dict)
      .repeat())

for x in ds:
  print(x)
  break

Выход: 

{'input_ids': <tf.Tensor: shape=(8,), dtype=int32, numpy=
array([  101, 13366,  2131,  1035,  6819,  2094,  1035,   102],
      dtype=int32)>, 'attention_mask': <tf.Tensor: shape=(8,), dtype=int32, numpy=array([1, 1, 1, 1, 1, 1, 1, 1], dtype=int32)>}
...