如何将TensorFlow MapDataset高效转换为tf.TensorArray？

问题背景 #

以下是生成TensorFlow MapDataset的示例代码：

import numpy as np
import tensorflow as tf

simple_data_samples = np.array([
         [1, 1, 1, -1, -1],
         [2, 2, 2, -2, -2],
         [3, 3, 3, -3, -3],
         [4, 4, 4, -4, -4],
         [5, 5, 5, -5, -5],
         [6, 6, 6, -6, -6],
         [7, 7, 7, -7, -7],
         [8, 8, 8, -8, -8],
         [9, 9, 9, -9, -9],
         [10, 10, 10, -10, -10],
         [11, 11, 11, -11, -11],
         [12, 12, 12, -12, -12],
])

def timeseries_dataset_multistep_combined(features, label_slice, input_sequence_length, output_sequence_length, batch_size):
    feature_ds = tf.keras.preprocessing.timeseries_dataset_from_array(features, None, input_sequence_length + output_sequence_length, batch_size=batch_size)

    def split_feature_label(x):
        x=tf.strings.as_string(x)

        return x[:, :input_sequence_length, :], x[:, input_sequence_length:, label_slice]

    feature_ds = feature_ds.map(split_feature_label)

    return feature_ds

ds = timeseries_dataset_multistep_combined(simple_data_samples, slice(None, None, None), input_sequence_length=4, output_sequence_length=2,
batch_size=1)
def print_dataset(ds):
    for inputs, targets in ds:
        print("---Batch---")
        print("Feature:", inputs.numpy())
        print("Label:", targets.numpy())
        print("")

print_dataset(ds)

变量ds为生成的TensorFlow MapDataset，请问将其转换为tf.TensorArray的最快且最高效的方法是什么？

解决方案 #

要实现最高效的转换，核心是利用TensorFlow的批量操作替代Python循环，具体步骤如下：

1. 先将数据集转为完整张量 #

由于数据集的每个元素是(inputs, targets)的结构，先分别收集所有批次的输入和目标，通过tf.concat堆叠成完整张量：

# 提取并拼接所有输入张量
all_inputs = tf.concat(list(ds.map(lambda x, y: x)), axis=0)
# 提取并拼接所有目标张量
all_targets = tf.concat(list(ds.map(lambda x, y: y)), axis=0)

2. 从完整张量初始化TensorArray #

使用tf.TensorArray.unstack方法直接从张量初始化TensorArray，这是底层优化的批量操作，比逐元素写入效率高得多：

# 转换输入为TensorArray
input_tensor_array = tf.TensorArray(dtype=all_inputs.dtype, size=tf.shape(all_inputs)[0])
input_tensor_array = input_tensor_array.unstack(all_inputs)

# 转换目标为TensorArray
target_tensor_array = tf.TensorArray(dtype=all_targets.dtype, size=tf.shape(all_targets)[0])
target_tensor_array = target_tensor_array.unstack(all_targets)

也可以写成更简洁的链式调用：

input_tensor_array = tf.TensorArray(dtype=all_inputs.dtype, size=tf.shape(all_inputs)[0]).unstack(all_inputs)
target_tensor_array = tf.TensorArray(dtype=all_targets.dtype, size=tf.shape(all_targets)[0]).unstack(all_targets)

效率说明 #

• 这种方法避免了Python层面的循环迭代，所有操作都在TensorFlow的图执行模式下完成，充分利用了TF的并行计算优化
• tf.concat和unstack都是批量操作，相比逐元素添加到TensorArray的方式，开销大幅降低
• 如果数据集规模极大导致内存不足，可以分批次处理，但上述方法是内存充足场景下的最优解

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