python - Tensorflow:分配变量后保存模型时出现 "GraphDef cannot be larger than 2GB."错误

Question

我想用一个预训练的模型来热启动另一个稍微有点不同的模型。简单地说，我创建了一个新模型，并为具有相同名称的变量分配预训练模型权重。但是，在保存模型时，出现了错误。


Traceback (most recent call last):
File "tf_test.py", line 23, in <module>
save_path = saver.save(sess, "./model.ckpt")
File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/training/saver.py", line 1308, in save
self.export_meta_graph(meta_graph_filename)
File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/training/saver.py", line 1331, in export_meta_graph
graph_def=ops.get_default_graph().as_graph_def(add_shapes=True),
File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/ops.py", line 2268, in as_graph_def
result, _ = self._as_graph_def(from_version, add_shapes)
File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/ops.py", line 2231, in _as_graph_def
raise ValueError("GraphDef cannot be larger than 2GB.")
ValueError: GraphDef cannot be larger than 2GB.


示例代码如下:

import tensorflow as tf
import numpy as np

v1 = tf.get_variable("L_enc", [400000, 1024])
v2 = tf.get_variable("L_dec", [400000, 1024])

init_op = tf.initialize_all_variables()

saver = tf.train.Saver(tf.all_variables())

with tf.Session() as sess:
  sess.run(init_op)
  for v in tf.trainable_variables():
    embedding = np.random.uniform(-1, 1, (400000, 1024))
    sess.run(v.assign(embedding))
  # Save the variables to disk.
  save_path = saver.save(sess, "./model.ckpt")
  print("Model saved in file: %s" % save_path)

Answer 1

法布里齐奥 correctly points out Protocol Buffer 的大小有 2GB 的硬性限制，但您可能想知道为什么您的程序会达到该限制。问题源于这些行:

for v in tf.trainable_variables():
  embedding = np.random.uniform(-1, 1, (400000, 1024))
  sess.run(v.assign(embedding))

当执行命中 v.assign(embedding) 时，新节点将添加到 TensorFlow 图中。特别是，每个 embedding 数组都转换为 tf.constant()张量，这将非常大(据我估计大约 328MB)。

避免这种情况的最佳方法是使用 tf.train.Saver 将先前模型中的变量直接加载到新模型中.由于模型可能具有不同的结构，您可能需要指定从旧模型中的变量名称到新模型中的 tf.Variable 对象的映射。

---

另一种解决问题的方法是预先创建一个 tf.placeholder() op 为每个变量赋值。这可能需要对您的实际代码进行更多重组，但以下内容对我有用:

v1 = tf.get_variable("L_enc", [400000, 1024])
v2 = tf.get_variable("L_dec", [400000, 1024])

# Define a separate placeholder and assign op for each variable, so
# that we can feed the initial value without adding it to the graph.
vars = [v1, v2]
placeholders = [tf.placeholder(tf.float32, shape=[400000, 1024]) for v in vars]
assign_ops = [v.assign(p) for (v, p) in zip(vars, placeholders)]

init_op = tf.global_variables_initializer()

saver = tf.train.Saver(tf.all_variables())

with tf.Session() as sess:
  sess.run(init_op)
  for p, assign_op in zip(placeholders, assign_ops):
    embedding = np.random.uniform(-1, 1, (400000, 1024))
    sess.run(assign_op, {p: embedding})

  # Save the variables to disk.
  save_path = saver.save(sess, "./model.ckpt")
  print("Model saved in file: %s" % save_path)