使用tf.data.Dataset.from_generator仅加载前256个元素并重复的问题

你遇到的问题核心是数据类型不匹配：你指定生成器输出的 dtype 为tf.uint8，但uint8的取值范围仅为0-255，当数值超过255时会自动执行模256运算，导致256被转为0、257转为1，以此类推，看起来像是重复前256个元素。

解决方法 #

将生成器的输出 dtype 替换为支持更大数值范围的整数类型，比如tf.int32或tf.int64，这两个类型完全能覆盖0到8999的数值区间。

修改后的代码：

def gen():
  for idx in z:
    yield idx

z = list(range(9000))  # 我的数据集长度为9000
dataset = tf.data.Dataset.from_generator(gen, tf.int32)  # 调整dtype类型

for step, sample in enumerate(dataset):
  print(step)  
  print(sample)

修改后遍历数据集时，会依次输出0到8999的对应张量，不会再出现循环重复的现象。

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