使用tf.data.Dataset时模型保存恢复后迭代器未初始化问题的解决方法

你遇到的这个错误非常典型——当你用tf.data.Dataset构建输入 pipeline 时，迭代器的状态不会被tf.train.Saver保存。训练时你初始化了迭代器，但恢复模型后，这个迭代器回到了未初始化的状态，而你的模型计算依赖于iter.get_next()的输出，自然就会报GetNext() failed because the iterator has not been initialized的错误。

先看你的测试代码问题：你直接导入图后就去运行preds:0，但此时图里的迭代器根本没被初始化，而且你的模型计算链路是placeholder -> dataset -> iterator -> x_inputs -> preds，直接给X:0喂数据没用，必须通过迭代器来流转数据。

解决方案步骤 #

1. 恢复模型后，从图中重新获取迭代器的初始化操作、相关占位符
2. 先执行迭代器的初始化，传入测试数据和对应的batch size
3. 再运行预测节点获取结果

修改后的完整代码 #

import tensorflow as tf
import os
import numpy as np
import math

# 生成模拟数据
features=np.random.randn(100,10)
features_test=np.random.randn(10,10)
y=np.random.randn(100,1)
y_test=np.random.randn(10,1)

feature_size=features.shape[1]
state_size=5
learning_rate=0.001

# 构建训练图
graph = tf.Graph()
with graph.as_default():
    batch_size_tensor = tf.placeholder(tf.int64,name="Batch_tensor")
    X,Y = tf.placeholder(tf.float32, [None,feature_size],"X"),tf.placeholder(tf.float32,[None,1],name="Y")
    
    dataset =tf.data.Dataset.from_tensor_slices((X,Y)).batch(batch_size_tensor).repeat()
    iter = dataset.make_initializable_iterator()
    x_inputs,y_outputs = iter.get_next()
    
    # 模型结构
    Wx = tf.Variable(tf.truncated_normal([feature_size, state_size], stddev=2.0 / math.sqrt(state_size)),name="Visual_weights_layer1")
    bx= tf.Variable(tf.zeros([state_size]),name="Visual_bias_layer1")
    x_hidden_state=tf.matmul(x_inputs, Wx)+bx
    x_hidden_state = tf.contrib.layers.batch_norm(x_hidden_state, epsilon=1e-5)
    vx=tf.nn.relu(x_hidden_state)
    
    W_final = tf.Variable(tf.truncated_normal([state_size, 1], stddev=2.0 / math.sqrt(state_size)),name="FinalLayer_weights")
    by=tf.Variable(tf.zeros([1]),name="FinalLayer_bias")
    predictions = tf.add(tf.matmul(vx, W_final), by,name="preds")
    
    loss = tf.losses.mean_squared_error(y_outputs,predictions)
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss)
    
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()

num_steps=100
batch_size=1
saver_path_model='tmp/testmodel'
export_path_model='tmp/testmodel.meta'

# 训练并保存模型
with tf.Session(graph=graph) as sess:
    sess.run(init)
    sess.run(iter.initializer, feed_dict={X: features, Y: y, batch_size_tensor: batch_size})
    print('initialized.')
    for step in range(num_steps):
        _, loss_val = sess.run([optimizer, loss])
        print(loss_val)
    saver.save(sess, saver_path_model)
    # 注：saver.save()会自动生成.meta文件，无需单独调用export_meta_graph

# 恢复模型并测试
sess = tf.Session()
new_saver = tf.train.import_meta_graph(export_path_model)
new_saver.restore(sess, saver_path_model)
graph = tf.get_default_graph()

# 从图中获取需要的节点
batch_size_tensor = graph.get_tensor_by_name("Batch_tensor:0")
X = graph.get_tensor_by_name("X:0")
Y = graph.get_tensor_by_name("Y:0")
iter_initializer = graph.get_operation_by_name("MakeIterator")
predictions_tensor = graph.get_tensor_by_name("preds:0")

# 初始化测试用的迭代器，这里batch_size设为测试集大小，一次性获取所有结果
test_batch_size = features_test.shape[0]
sess.run(iter_initializer, feed_dict={X: features_test, Y: y_test, batch_size_tensor: test_batch_size})

# 获取预测结果
predictions_test = sess.run(predictions_tensor)
print("测试集预测结果：", predictions_test)

关键注意点 #

• 迭代器状态不保存：tf.train.Saver只会保存模型的变量（比如权重、偏置），迭代器的状态属于图的临时操作状态，不会被保存，所以每次加载模型后必须重新初始化。
• 数据路径一致：测试时的数据必须和训练时一样通过dataset -> iterator的路径流入模型，不能直接给原始placeholder喂数据（因为模型计算依赖的是迭代器输出的x_inputs，不是原始placeholder）。

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