使用map方法构建TensorFlow Dataset并应用TextVectorization层报错排查

我正在为NLP模型的Embedding层构建输入，需将原始文本转换为Embedding层所需的数值型数据，过程中遇到如下问题：

示例数据 #

# 0 = negative
# 1 = positive
documents = [['topology freaking sucks man, what a waste of time!', 0], ['wow bro you a NLP fan? Tell me more I want to know', 1], 
['you know, I will eventually die',0], ['the secret to happiness is to only be depresssed',0], 
['what is the floor without feet', 1], ['regicide is permissable only in historical situations',1],
['I do not like delivering wehat based products for I am allergic to wheat', 0], 
['Why does he ring the large bell every hour?',0],
['Wisdom comes not from experience but from knowing',1], 
['Little is known of the inner workings of the feline mind', 1]]

每条数据含一个句子和正负标签，符合文档分类任务要求。

构建TextVectorization层 #

# create preprocessing layer
VOCAB_SIZE = 500 # max amount of vocabulary amongst all documents
MAX_SEQUENCE_LENGTH = 50 # maximum amount of words/tokens that will be considered in each document
# output mode 'int' will assign unique integer per token, so in our example below, 'topology' is assigned the value
# 19. Notice that these integers are randomly assigned and essentially acts as a hashmap
int_vectorize_layer = TextVectorization(
    max_tokens=VOCAB_SIZE,
    output_mode='int',
    output_sequence_length = MAX_SEQUENCE_LENGTH
)

数据集构建与报错情况 #

将原始数据拆分转换为Tensor后，通过生成器构建TensorFlow Dataset，适配TextVectorization层后，调用Dataset的map方法应用int_vectorize_text函数时出现TypeError（提示'NoneType'与'int'无法比较），但手动遍历Dataset调用该函数却能正常输出结果。

相关代码 #

# Applies adapted layer to tensorflow dataset
def int_vectorize_text(sentence, label):
  sentence = tf.expand_dims(sentence, -1)
  sentence = tf.squeeze(sentence, axis=-1)
  return int_vectorize_layer(sentence), label


# passes raw data as a generator to the Dataset from_generator constructor
def generate_data(sentences, labels):
  for s, l in zip(sentences,labels):
    yield s, l

# split raw data between training and validation set
train_docs = documents[:8]
val_docs = documents[8:]

# separate sentences and labels
train_sentences = [d[0] for d in train_docs]
train_labels = [d[1] for d in train_docs]

val_sentences = [d[0] for d in val_docs]
val_labels = [d[1] for d in val_docs]

# convert to tensors
train_sentences_tensor = tf.convert_to_tensor(train_sentences)
train_labels_tensor = tf.convert_to_tensor(train_labels)

val_sentences_tensor = tf.convert_to_tensor(val_sentences)
val_labels_tensor = tf.convert_to_tensor(val_labels)

# build tensorflow Dataset using the above generator function on the newly constructed tensor objects
train_dataset = tf.data.Dataset.from_generator(
    generate_data, (tf.string, tf.int32), args=(train_sentences_tensor, train_labels_tensor))
val_dataset = tf.data.Dataset.from_generator(
    generate_data, (tf.string, tf.int32), args=(val_sentences_tensor, val_labels_tensor))

# adapt layer using training sentences
int_vectorize_layer.adapt(train_sentences)

# now here is where the error occurs
int_train_df = train_dataset.map(int_vectorize_text) # ERROR
int_val_df = val_dataset.map(int_vectorize_text)

错误信息 #

TypeError                                 Traceback (most recent call last)
/home/akagi/Documents/Projects/MLMastery NLP Tutorial/Lesson 5 - Learned Embedding.ipynb Cell 7 in <cell line: 21>()
     19 # Use the map method to apply the int_vectorize_text function to each element of the dataset
     20 int_vectorize_layer.adapt(train_sentences)
---> 21 int_train_df = train_dataset.map(int_vectorize_text)
     22 int_val_df = val_dataset.map(int_vectorize_text)

File ~/Documents/Projects/.venv/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py:2294, in DatasetV2.map(self, map_func, num_parallel_calls, deterministic, name)
   2291   if deterministic is not None and not DEBUG_MODE:
   2292     warnings.warn("The `deterministic` argument has no effect unless the "
   2293                   "`num_parallel_calls` argument is specified.")
-> 2294   return MapDataset(self, map_func, preserve_cardinality=True, name=name)
   2295 else:
   2296   return ParallelMapDataset(
   2297       self,
   2298       map_func,
   (...)
   2301       preserve_cardinality=True,
   2302       name=name)

File ~/Documents/Projects/.venv/lib/python3.8/site-packages/tensorflow/python/data/ops/dataset_ops.py:5499, in MapDataset.__init__(self, input_dataset, map_func, use_inter_op_parallelism, preserve_cardinality, use_legacy_function, name)
   5497 self._use_inter_op_parallelism = use_inter_op_parallelism
   5498 self._preserve_cardinality = preserve_cardinality
-> 5499 self._map_func = structured_function.StructuredFunctionWrapper(
...
    '>' not supported between instances of 'NoneType' and 'int'
    
    Call arguments received by layer 'text_vectorization' (type TextVectorization):
      • inputs=tf.Tensor(shape=<unknown>, dtype=string)

Dataset结构验证 #

(<tf.Tensor: shape=(), dtype=string, numpy=b'topology freaking sucks man, what a waste of time!'>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'wow bro you a NLP fan? Tell me more I want to know'>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'you know, I will eventually die'>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'the secret to happiness is to only be depresssed'>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'what is the floor without feet'>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'regicide is permissable only in historical situations'>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'I do not like delivering wehat based products for I am allergic to wheat'>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)
(<tf.Tensor: shape=(), dtype=string, numpy=b'Why does he ring the large bell every hour?'>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)

手动调用正常 #

for x in train_dataset:
    print(int_vectorize_text(x[0], x[1]))

请问该报错的原因是什么？如何解决？

报错核心原因：from_generator创建的Dataset未明确元素形状信息，map操作触发TensorFlow图模式编译时，TextVectorization层无法处理形状未知的输入张量，导致内部逻辑出现NoneType与整数的比较，触发TypeError。而手动遍历是在Eager模式下执行，TensorFlow能动态获取张量形状，因此不会报错。

另外，int_vectorize_text函数中的expand_dims和squeeze操作完全冗余，既未改变有效数据，还可能干扰形状推断。

推荐两种修改方式，优先选择方式一：

方式一：替换from_generator为更高效的Dataset创建方式 #

既然已将数据转换为张量，无需使用生成器，直接用tf.data.Dataset.from_tensor_slices创建Dataset，它会自动保留元素形状信息：

# 替换原Dataset创建代码
train_dataset = tf.data.Dataset.from_tensor_slices((train_sentences_tensor, train_labels_tensor))
val_dataset = tf.data.Dataset.from_tensor_slices((val_sentences_tensor, val_labels_tensor))

# 简化处理函数
def int_vectorize_text(sentence, label):
  return int_vectorize_layer(sentence), label

# 后续操作保持不变
int_vectorize_layer.adapt(train_sentences)
int_train_df = train_dataset.map(int_vectorize_text)
int_val_df = val_dataset.map(int_vectorize_text)

方式二：为生成器输出指定明确形状（保留生成器场景） #

若必须使用from_generator，需在生成器中明确指定张量形状，并修改from_generator的输出签名：

def generate_data(sentences, labels):
  for s, l in zip(sentences,labels):
    # 明确指定标量形状
    yield tf.constant(s, shape=()), tf.constant(l, shape=())

# 修改from_generator的输出签名
train_dataset = tf.data.Dataset.from_generator(
    generate_data, 
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.string),
        tf.TensorSpec(shape=(), dtype=tf.int32)
    ),
    args=(train_sentences_tensor, train_labels_tensor))

val_dataset = tf.data.Dataset.from_generator(
    generate_data, 
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.string),
        tf.TensorSpec(shape=(), dtype=tf.int32)
    ),
    args=(val_sentences_tensor, val_labels_tensor))

# 简化处理函数
def int_vectorize_text(sentence, label):
  return int_vectorize_layer(sentence), label

# 后续操作不变
int_vectorize_layer.adapt(train_sentences)
int_train_df = train_dataset.map(int_vectorize_text)
int_val_df = val_dataset.map(int_vectorize_text)

内容的提问来源于stack exchange，提问作者hkj447