使用tf.data.Dataset导致TensorFlow保存模型体积过大问题咨询

我来帮你捋清楚这个问题——在TensorFlow 1.4里从旧的队列输入（tf.train.string_input_producer()+tf.train.batch()）切换到tf.data.Dataset后，模型体积明显变大，核心原因和两者的状态存储机制差异有关：

为什么体积会变大？ #

• 旧队列机制的状态是临时的：用tf.train.start_queue_runners()启动的队列线程，其运行状态（比如队列里的待处理数据、指针位置）是存在会话的临时线程中，不属于TensorFlow的全局变量，所以调用saver.save()时不会把这些状态保存下来，模型文件只包含网络的可训练参数（权重、偏置等）。
• tf.data.Dataset的状态会被持久化：在TF1.x早期版本（比如1.4）中，tf.data的迭代器、缓存缓冲、shuffle的缓冲区状态等都会被注册为全局变量。当你保存模型时，这些和数据管道相关的状态变量会被一起保存，直接导致模型文件体积膨胀。

具体解决办法 #

1. 精准指定要保存的变量列表（最推荐） #

创建Saver时，只传入模型的可训练参数，排除所有和tf.data相关的变量：

# 方案一：只保存可训练变量（网络权重、偏置等）
trainable_vars = tf.trainable_variables()
saver = tf.train.Saver(var_list=trainable_vars)

# 方案二：更精准过滤，排除tf.data相关变量
vars_to_save = [var for var in tf.global_variables() 
                if not any(key in var.name for key in ["Iterator", "dataset", "batch", "prefetch", "shuffle"])]
saver = tf.train.Saver(var_list=vars_to_save)

这样保存的模型就只会包含训练所需的核心参数，和旧队列机制下的体积一致。

2. 优化tf.data的迭代器使用 #

在TF1.4中，尽量使用one-shot迭代器（不需要初始化的迭代器），它的状态变量更少，相比可初始化迭代器或重新初始化迭代器，能减少不必要的状态存储：

# 示例：使用one-shot迭代器
dataset = tf.data.TextLineDataset("image_paths.txt")
# 后续处理：解析图像、batch等
iterator = dataset.make_one_shot_iterator()
next_batch = iterator.get_next()

不过这个方法的效果不如第一种明显，因为还是会有一些基础状态变量存在。

3. 避免数据管道中的冗余状态 #

如果你的tf.data pipeline里用了cache()、shuffle(buffer_size=大数值)这类操作，会在内存中缓存大量数据或状态，这些都会被保存到模型里。如果不需要持久化这些状态，尽量在训练结束前清理，或者调整参数减少缓存量。

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