使用tf.data.Dataset时tf.train.write_graph引发内存错误的原因探究

为什么tf.data.Dataset会导致导出计算图时内存溢出？ #

在TensorFlow 1.x的静态图模式下，tf.data.Dataset的所有构建操作（比如from_tensor_slices、shuffle、batch、make_one_shot_iterator）都是在图构建阶段完成的，这会带来一个关键问题：当你用from_tensor_slices传入大的numpy数组（比如MNIST的60000条训练数据）时，TensorFlow会把整个数组作为常量节点直接嵌入到计算图中。

这就意味着你的计算图不仅包含模型的层、损失、优化器这些逻辑节点，还硬生生塞了几十MB甚至更大的训练数据。当你执行tf.train.write_graph时，需要把整个图（包括嵌在里面的数据集）导出成文本格式的Proto Buffer，这会瞬间占用大量内存——哪怕你的模型本身很小，也会因为嵌在图里的数据集导致内存爆掉。

而你之前用Keras或者纯numpy feed_dict的方式时，数据是在运行时通过feed_dict传入的，并没有被存在计算图里，所以计算图本身只有模型结构，体积很小，导出完全没问题。

解决办法 #

针对这个问题，有几个可行的修复方案：

1. 避免直接将大numpy数组嵌入计算图

改用tf.data.Dataset.from_generator或者从文件（比如TFRecord）读取数据，这样计算图里只会保留数据生成/读取的逻辑，不会包含实际的数据集：

def data_generator():
    for x, y in zip(x_train, y_train):
        yield x, y

# 用生成器构建数据集
data_pipeline = tf.data.Dataset.from_generator(
    data_generator,
    output_types=(tf.float32, tf.float32),
    output_shapes=(tf.TensorShape([784]), tf.TensorShape([10]))
)
iter = data_pipeline.shuffle(1000).repeat().batch(1024).make_one_shot_iterator()
next_item = iter.get_next()

2. 单独导出模型结构的子图

如果你一定要用from_tensor_slices，可以在导出时只保留模型相关的节点，排除数据管道的部分。比如先定义好模型的输入输出张量，然后导出以这些张量为核心的子图：

# 定义模型的输入输出
model_input = X
model_output = logits

# 只导出模型相关的子图
subgraph = tf.graph_util.extract_sub_graph(sess.graph_def, [model_output.name.split(':')[0]])
tf.train.write_graph(subgraph, './', 'model_only.pbtxt', as_text=True)

3. 升级到TensorFlow 2.x

TF2.x默认采用即时执行模式，计算图是动态构建的，tf.data的实现逻辑也做了优化，不会把数据集常量嵌入到图中。同时TF2.x的模型导出（SavedModel）更稳定，也更符合现代TensorFlow的使用习惯。

tf.data.Dataset的使用建议 #

并不是说tf.data只能在大数据或高性能机器上用——它的优势（并行读取、预取、集成数据增强等）在小数据集上也能体现，只是在TF1.x里要注意不要直接把大numpy数组通过from_tensor_slices塞进计算图。

对于小数据集，用feed_dict确实更简单省心；但如果需要用到tf.data的高级功能，只要避免把数据集常量嵌进图里，普通8GB内存的机器完全可以正常运行。

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