基于tf.contrib.data.prefetch_to_device，如何仅预取小批量训练数据而非标签？

刚好碰到过类似的场景，针对你提到的tf.nn.ctc_loss这类运行在CPU上的损失函数，我们可以通过拆分数据集元素的处理逻辑，让tf.contrib.data.prefetch_to_device只作用于特征部分，具体实现步骤如下：

核心思路 #

tf.contrib.data.prefetch_to_device默认会对整个数据集元素（比如包含特征和标签的元组）进行预取，但我们可以通过dataset.map()操作，单独对特征应用预取逻辑，标签则保持在CPU内存中，完美适配CPU端损失函数的需求。

代码示例 #

import tensorflow as tf

# 1. 先构建你的原始数据集，每个元素是(features, labels)的元组
# 这里只是示例，替换成你实际的数据集构建代码
raw_dataset = tf.data.Dataset.from_generator(
    your_data_generator,
    output_types=(tf.float32, tf.int32),
    output_shapes=((None, 128), (None,))
)

# 2. 指定要预取到的GPU设备
target_gpu = '/GPU:0'

# 3. 定义处理函数：仅将特征预取到GPU，标签留在CPU
def prefetch_features(features, labels):
    # 对特征应用预取到设备的操作
    prefetched_features = tf.contrib.data.prefetch_to_device(target_gpu)(features)
    # 返回处理后的特征和原标签
    return prefetched_features, labels

# 4. 将处理函数应用到数据集
processed_dataset = raw_dataset.map(prefetch_features)

# 5. 后续可以继续添加batch、shuffle等操作
processed_dataset = processed_dataset.shuffle(1000).batch(32)

关键说明 #

• 经过map处理后，每个数据集元素中的特征张量已经被预取到指定GPU，而标签张量仍保留在CPU内存中，完全符合tf.nn.ctc_loss不需要标签在GPU的要求。
• 如果你的TensorFlow版本是2.x及以上，tf.contrib模块已经被移除，对应的替代API是tf.data.experimental.prefetch_to_device，用法和上述示例一致，只需要替换函数名即可。

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