如何使用TensorFlow Dataset API实现训练样本重复复制（含不同次数场景）

一、统一重复每个样本指定次数 #

tf.data.Dataset.repeat()是对整个数据集做重复循环，要实现单个样本重复指定次数，可结合flat_map和子数据集生成来实现：

import tensorflow as tf

# 初始化原始数据集
dataset = tf.data.Dataset.from_tensor_slices([1,2,3,4,5])
# 设定每个样本的重复次数
repeat_times = 3

# 核心逻辑：对每个样本生成包含repeat_times个自身的子数据集，再扁平化合并
dataset = dataset.flat_map(lambda x: tf.data.Dataset.from_tensor_slices([x] * repeat_times))

# 输出验证
for elem in dataset:
    print(elem.numpy(), end=' ')
# 结果：1 1 1 2 2 2 3 3 3 4 4 4 5 5 5

也可以用tf.repeat配合map实现，性能更优：

dataset = dataset.map(lambda x: tf.repeat(x, repeat_times)).flat_map(lambda x: tf.data.Dataset.from_tensor_slices(x))

二、按权重数组指定每个样本的重复次数 #

如果每个样本的重复次数由weights数组单独指定（比如weights = [2,3,1,4,2]，对应样本1重复2次、样本2重复3次），可以先将样本与对应权重配对，再逐对处理：

import tensorflow as tf

samples = [1,2,3,4,5]
weights = [2,3,1,4,2]

# 将样本和对应权重配对成数据集
dataset = tf.data.Dataset.from_tensor_slices((samples, weights))
# 对每个(样本, 重复次数)对，生成对应次数的样本序列并扁平化
dataset = dataset.flat_map(lambda x, cnt: tf.data.Dataset.from_tensor_slices(tf.repeat(x, cnt)))

# 输出验证
for elem in dataset:
    print(elem.numpy(), end=' ')
# 结果：1 1 2 2 2 3 4 4 4 4 5 5

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