Эффективное преобразование токенов в векторы слов с помощью TensorFlow Transform

Question

Я хотел бы использовать TensorFlow Transform для преобразования токенов в векторы слов на этапе обучения, проверки и вывода.

Я последовал за этим постом Stackru и осуществил первоначальное преобразование из токенов в векторы. Преобразование работает, как ожидалось, и я получаю векторы EMB_DIM за каждый токен.

import numpy as np
import tensorflow as tf

tf.reset_default_graph()
EMB_DIM = 10

def load_pretrained_glove():
    tokens = ["a", "cat", "plays", "piano"]
    return tokens, np.random.rand(len(tokens), EMB_DIM)

# sample string 
string_tensor = tf.constant(["plays", "piano", "unknown_token", "another_unknown_token"])


pretrained_vocab, pretrained_embs = load_pretrained_glove()

vocab_lookup = tf.contrib.lookup.index_table_from_tensor(
    mapping = tf.constant(pretrained_vocab),
    default_value = len(pretrained_vocab))
string_tensor = vocab_lookup.lookup(string_tensor)

# define the word embedding
pretrained_embs = tf.get_variable(
    name="embs_pretrained",
    initializer=tf.constant_initializer(np.asarray(pretrained_embs), dtype=tf.float32),
    shape=pretrained_embs.shape,
    trainable=False)

unk_embedding = tf.get_variable(
    name="unk_embedding",
    shape=[1, EMB_DIM],
    initializer=tf.random_uniform_initializer(-0.04, 0.04),
    trainable=False)

embeddings = tf.cast(tf.concat([pretrained_embs, unk_embedding], axis=0), tf.float32)
word_vectors = tf.nn.embedding_lookup(embeddings, string_tensor)

with tf.Session() as sess:
    tf.tables_initializer().run()
    tf.global_variables_initializer().run()
    print(sess.run(word_vectors))

Когда я рефакторинг кода для запуска в виде TFX Transform Graph, я получаю сообщение об ошибке ConversionError ниже.

import pprint
import tempfile
import numpy as np
import tensorflow as tf
import tensorflow_transform as tft
import tensorflow_transform.beam.impl as beam_impl
from tensorflow_transform.tf_metadata import dataset_metadata
from tensorflow_transform.tf_metadata import dataset_schema

tf.reset_default_graph()

EMB_DIM = 10

def load_pretrained_glove():
    tokens = ["a", "cat", "plays", "piano"]
    return tokens, np.random.rand(len(tokens), EMB_DIM)


def embed_tensor(string_tensor, trainable=False):
    """
    Convert List of strings into list of indices then into EMB_DIM vectors
    """

    pretrained_vocab, pretrained_embs = load_pretrained_glove()

    vocab_lookup = tf.contrib.lookup.index_table_from_tensor(
        mapping=tf.constant(pretrained_vocab),
        default_value=len(pretrained_vocab))
    string_tensor = vocab_lookup.lookup(string_tensor)

    pretrained_embs = tf.get_variable(
        name="embs_pretrained",
        initializer=tf.constant_initializer(np.asarray(pretrained_embs), dtype=tf.float32),
        shape=pretrained_embs.shape,
        trainable=trainable)
    unk_embedding = tf.get_variable(
        name="unk_embedding",
        shape=[1, EMB_DIM],
        initializer=tf.random_uniform_initializer(-0.04, 0.04),
        trainable=False)

    embeddings = tf.cast(tf.concat([pretrained_embs, unk_embedding], axis=0), tf.float32)
    return tf.nn.embedding_lookup(embeddings, string_tensor)

def preprocessing_fn(inputs):
    input_string = tf.string_split(inputs['sentence'], delimiter=" ") 
    return {'word_vectors': tft.apply_function(embed_tensor, input_string)}


raw_data = [{'sentence': 'This is a sample sentence'},]
raw_data_metadata = dataset_metadata.DatasetMetadata(dataset_schema.Schema({
  'sentence': dataset_schema.ColumnSchema(
      tf.string, [], dataset_schema.FixedColumnRepresentation())
}))

with beam_impl.Context(temp_dir=tempfile.mkdtemp()):
    transformed_dataset, transform_fn = (  # pylint: disable=unused-variable
        (raw_data, raw_data_metadata) | beam_impl.AnalyzeAndTransformDataset(
            preprocessing_fn))

    transformed_data, transformed_metadata = transformed_dataset  # pylint: disable=unused-variable
    pprint.pprint(transformed_data)

Сообщение об ошибке

TypeError: Failed to convert object of type <class 
'tensorflow.python.framework.sparse_tensor.SparseTensor'> to Tensor. 
Contents: SparseTensor(indices=Tensor("StringSplit:0", shape=(?, 2), 
dtype=int64), values=Tensor("hash_table_Lookup:0", shape=(?,), 
dtype=int64), dense_shape=Tensor("StringSplit:2", shape=(2,), 
dtype=int64)). Consider casting elements to a supported type.

Вопросы

1. Почему шаг TF Transform требует дополнительного преобразования / приведения?
2. Возможен ли такой подход преобразования токенов в векторы слов? Векторы слова могут быть несколько гигабайт в памяти. Как Apache Beam работает с векторами? Если Beam в распределенной установке, потребует ли это N x vector memory с N количество рабочих?

Answer 1

Ошибка, связанная с SparseTensor, связана с тем, что вы вызываете string_split, который возвращает SparseTensor. Ваш тестовый код не вызывает string_split, поэтому это происходит только с вашим кодом Transform.

Что касается памяти, вы правы, матрица внедрения должна быть загружена в каждого работника.

Answer 2

Нельзя поместить SparseTensor в словарь, возвращаемый преобразованием TFX, в вашем случае функцией preprocessing_fn. Причина в том, что SparseTensor не является Tensor, это на самом деле небольшой подграф.

Чтобы исправить ваш код, вы можете преобразовать SparseTensor в Tensor. Существует несколько способов сделать это, я бы рекомендовал использовать tf.serialize_sparse для обычного SparseTensor и tf.serialize_many_sparse для пакетного.

Чтобы использовать такой сериализованный Tensor в Trainer, вы можете вызвать функцию tf. deserialize_many_sparse.