TensorFlow程序第5次数据检索时触发StopIteration错误

你遇到的问题是：运行一段基于TensorFlow Eager执行模式的代码处理iris训练数据集，当指定循环5次获取batch数据时，第5次调用iter.next()会抛出StopIteration错误。先看看你的代码和报错信息：

你的代码 #

#!/usr/bin/env python3
from __future__ import absolute_import, division, print_function
import sys
import tensorflow as tf
from utils import parse_csv

tf.enable_eager_execution()

train_dataset = tf.data.TextLineDataset("iris_training.csv")
# skip the first header row
train_dataset = train_dataset.skip(1)
# parse each row
train_dataset = train_dataset.map(parse_csv)
# randomize
train_dataset = train_dataset.shuffle(buffer_size=1000)
# batch
train_dataset = train_dataset.batch(32)

# iterator
iter = iter(train_dataset)

# View few entries from the batch
print()
for i in range(0, int(sys.argv[1])):
    features, label = iter.next()
    print("=> features:", features[0])
    print("=> label:", label[0])
    print()

报错信息 #

./tensor.py 5 2>&1 | tail -5
File "./tensor.py", line 33, in <module>
features, label = iter.next()
File "/Users/nv/.pyenv/versions/3.6.5/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 519, in next
raise StopIteration
StopIteration

原因分析 #

这个问题的核心是你的数据集总batch数不足5个，当迭代器耗尽所有数据后，再次调用next()就会触发Python迭代器的标准StopIteration异常。

具体来说：

• 标准的iris_training.csv包含120条训练数据（跳过表头后）
• 你设置的batch size是32，计算一下总batch数：
  • 32 × 3 = 96条数据，对应3个完整batch
  • 剩下120 - 96 = 24条数据，组成第4个不完整的batch
• 所以整个数据集总共只有4个batch，当你循环5次时，前4次能正常获取数据，第5次迭代器已经没有数据可以返回，就会抛出StopIteration。

解决方案 #

根据你的需求，有几种处理方式：

1. 用for循环自动遍历（推荐）

Python的for循环会自动处理迭代器的StopIteration，不需要手动调用next()，更简洁安全：

# 替换原来的循环部分
print()
for features, label in train_dataset:
    print("=> features:", features[0])
    print("=> label:", label[0])
    print()

2. 捕获异常避免崩溃

如果一定要固定次数循环，可以用try-except捕获异常，优雅处理数据耗尽的情况：

print()
for i in range(0, int(sys.argv[1])):
    try:
        features, label = iter.next()
        print("=> features:", features[0])
        print("=> label:", label[0])
        print()
    except StopIteration:
        print(f"注意：第{i+1}次迭代时已无更多数据")
        break

3. 让数据集循环迭代（适合训练场景）

如果是训练时需要多轮迭代数据，可以给数据集添加repeat()方法，这样迭代器会循环生成数据，不会轻易耗尽：

# 在batch之后添加repeat()
train_dataset = train_dataset.batch(32).repeat()

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