如何将tf.data.Dataset.map的结果返回为单个张量

你当前用numpy迭代器转列表再转张量的方法确实繁琐，而且不是TensorFlow推荐的方案，这里有两种更直接高效的方法：

• 方法一：批量打包后直接转张量
  利用tf.data.Dataset.batch()将整个数据集合并为一个批次，再用tf.convert_to_tensor()转换，这是最简洁的方式：

  # 假设map处理后的数据集为map_dataset
  total_samples = map_dataset.cardinality()
  single_tensor = tf.convert_to_tensor(map_dataset.batch(total_samples))
  

  cardinality()会返回数据集的样本总数，batch后所有数据会被打包成一个大张量，直接转换即可。

• 方法二：用reduce操作拼接元素
  如果需要自定义拼接逻辑，可以用tf.data.Dataset.reduce()逐步拼接所有元素：

  # 假设每个元素形状为(20, 2)，初始化空张量作为初始状态
  initial_tensor = tf.zeros((0, 20, 2), dtype=tf.float32)
  single_tensor = map_dataset.reduce(initial_tensor, lambda acc, elem: tf.concat([acc, elem], axis=0))
  

  这个方法会遍历数据集，把每个元素依次拼接到累加器中，最终得到完整的单一张量。

为什么不推荐原方法？ #

原方法需要将数据从TensorFlow的计算图环境转到Python环境，不仅步骤冗余，还可能降低大数据集的处理性能，而且如果在图模式下运行时，numpy相关操作会直接报错，因为图模式不允许这类eager模式的操作。

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