TensorFlow中Dataset.map运行自定义函数时张量shape[0]为None问题

核心原因是TensorFlow存在静态shape、动态shape两套取值逻辑，再加上Dataset.map默认运行在图模式：

• 直接迭代Dataset取元素时处于Eager执行模式，张量持有实际运行时值，通过tensor.shape[0]能直接拿到Python原生整数类型的长度值，计算不会出错。
• 给函数加@tf.function注解后在map外部调用时，传入的张量静态shape是确定的，tf.function做函数追踪时能拿到具体长度值，因此也能正常运行。
• 但在Dataset.map的执行流中，输入是变长的Ragged张量，每个元素的长度不固定，图构建阶段无法推导出固定的第0维长度，此时通过tensor.shape访问的是编译期静态shape属性，第0维返回值为None，直接拿None做整数除法自然触发类型错误。

图模式下处理动态长度张量时，不能依赖静态shape属性做运行时数值计算，必须通过TensorFlow内置算子获取运行时的动态值。

方案1：修正原有函数逻辑 #

把所有和张量长度相关的计算替换为TensorFlow内置算子，让计算逻辑嵌入图的执行流，在运行时动态获取每个元素的实际长度：

def batches_of_four(tokens):
  # 用tf.shape获取运行时的动态长度，返回可在图中计算的张量值
  token_length = tf.shape(tokens)[0]
  splits = token_length // 4
  tokens = tokens[:splits * 4]
  return tf.split(tokens, num_or_size_splits=splits)

修正后的函数可以直接传入Dataset.map正常运行，不会再出现None值计算的错误。

方案2：使用tf.data原生API实现（性能更优） #

不需要自定义切分函数，直接利用tf.data内置的流水线算子完成切分，框架会自动做执行优化，同时drop_remainder参数原生支持丢弃尾部不足长度的片段：

dataset = tf.data.Dataset.from_tensor_slices(
    tf.ragged.constant([[1, 2, 3, 4, 5], [4, 5, 6, 7]]))
# 对每个变长序列做切分，自动丢弃不足4个元素的尾部块
batched_dataset = dataset.flat_map(
    lambda seq: tf.data.Dataset.from_tensor_slices(seq).batch(4, drop_remainder=True)
)

运行后可以得到两个长度为4的张量：[1,2,3,4]和[4,5,6,7]，和预期逻辑完全一致。

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