在内存中序列化和反序列化TensorFlow模型并继续训练

Question

我看到过类似的问题，但是我还没有找到一个令人满意的答案。基本上，我想要在tensorflow中实现与keras中model.to_json()、model.get_weights()、model.from_json()和model.set_weights()相当的功能。我认为我已经接近目标了，但是我卡在了某个点上。如果可能的话，我更希望能够在同一个字符串中获取权重和图形，但我理解如果这不可能。

目前，我拥有的是：

g = optimizer.minimize(loss_op,
                       global_step=tf.train.get_global_step())
de = g.graph.as_graph_def()
json_string = json_format.MessageToJson(de)

gd = tf.GraphDef()
gd = json_format.Parse(json_string, gd)

看起来创建图表很不错，但显然变量、权重等元图并未包含在内。还有元图，但我只能看到export_meta_graph一项，它似乎不能以同样的方式序列化。我看到MetaGraph有一个proto函数，但我不知道如何序列化这些变量。

简而言之，您会如何将tensorflow模型（即权重、图形等）序列化为字符串（最好是json），然后反序列化并继续训练或提供预测？

以下是让我接近目标并尝试过的内容，但大多数需要写入磁盘，这在这种情况下是不可能的：

- Derek_M

Answer 1

请注意，@Maxim提供的解决方案每次运行都会在图中创建新的操作。

如果您经常运行此函数，这将导致代码变得越来越慢。

解决此问题的两种方法：

1. 同时创建赋值操作，与图的其余部分一起重用它们：

   assign_ops = []
   for var_name in tf.trainable_variables():
       assign_placeholder = tf.placeholder(var.dtype, shape=value.shape)
       assign_op = var.assign(assign_placeholder)
       assign_ops.append(assign_op)
   

2. 使用变量上的加载函数，我更喜欢这种方法，因为它不需要上面的代码：

   self.params = tf.trainable_variables()
   
   def get_weights(self):
       values = tf.get_default_session().run(self.params)
       return values
   
   def set_weights(self, weights):
       for i, value in enumerate(weights):
           value = np.asarray(value)
           self.params[i].load(value, self.sess)
   

（我不能评论，所以我把这个作为答案）

- Sebastien Kerbrat

Answer 2

您可以使用 freeze_graph 这个脚本，它包含在TensorFlow中，可以将GraphDef proto、SaverDef proto和存储在checkpoint文件中的一组变量值作为输入。

通过这种方式，您可以输出一个GraphDef，其中所有变量操作都转换为包含变量值的const操作。

要恢复冻结的模型，您需要重新初始化图表并从冻结的模型中重新映射输入，请参见此示例。

- Silverstorm

---

没错，那就是我链接的要点。它似乎仍然需要去磁盘，对吧？而且看起来那个要点没有使用tensorflow的冻结图函数（已注释掉），我很好奇是否有原因。 - Derek_M

Answer 3

如果你想要等价于keras Model.get_weights()和Model.set_weights()的方法，这些方法与keras内部并不紧密相连，可以很容易地提取出来。

原始代码

以下是它们在keras源代码中的样子：

def get_weights(self):
  weights = []
  for layer in self.layers:
    weights += layer.weights
  return K.batch_get_value(weights)   # this is just `get_session().run(weights)`

def set_weights(self, weights):
  tuples = []
  for layer in self.layers:
    num_param = len(layer.weights)
    layer_weights = weights[:num_param]
    for sw, w in zip(layer.weights, layer_weights):
      tuples.append((sw, w))
    weights = weights[num_param:]
  K.batch_set_value(tuples)  # another wrapper over `get_session().run(...)`

Keras的weights是numpy数组列表（不是json）。正如您所看到的，它利用了已知模型架构（self.layers）的事实，从而使其能够重建从变量到值的正确映射。一些看似非常复杂的工作在K.batch_set_value中完成，但实际上它只是准备分配操作并在会话中运行它们。

在纯TensorFlow中获取和设置权重

def tensorflow_get_weights():
  vars = tf.trainable_variables()
  values = tf.get_default_session().run(vars)
  return zip([var.name for var in vars], values)

def tensorflow_set_weights(weights):
  assign_ops = []
  feed_dict = {}
  for var_name, value in weights:
    var = tf.get_default_session().graph.get_tensor_by_name(var_name)
    value = np.asarray(value)
    assign_placeholder = tf.placeholder(var.dtype, shape=value.shape)
    assign_op = tf.assign(var, assign_placeholder)
    assign_ops.append(assign_op)
    feed_dict[assign_placeholder] = value
  tf.get_default_session().run(assign_ops, feed_dict=feed_dict)

我假设您希望在默认会话中序列化/反序列化整个模型（即所有可训练变量）。如果不是这种情况，上面的函数很容易进行自定义。

测试

x = tf.placeholder(shape=[None, 5], dtype=tf.float32, name='x')
W = tf.Variable(np.zeros([5, 5]), dtype=tf.float32, name='W')
b = tf.Variable(np.zeros([5]), dtype=tf.float32, name='b')
y = tf.add(tf.matmul(x, W), b)

with tf.Session() as session:
  session.run(tf.global_variables_initializer())

  # Save the weights
  w = tensorflow_get_weights()
  print(W.eval(), b.eval())

  # Update the model
  session.run([tf.assign(W, np.ones([5, 5])), tf.assign(b, np.ones([5]) * 2)])
  print(W.eval(), b.eval())

  # Restore the weights
  tensorflow_set_weights(w)
  print(W.eval(), b.eval())

如果您运行此测试，您应该看到模型被冻结在零处，然后得到更新，最后恢复为零。

- Maxim

---

太棒了。我今天早上会看一下。 - Derek_M

@Derek_M 当然可以。如果您有进一步的问题，请告诉我。 - Maxim

好的。我实际上去掉了变量，并将numpy权重转换为列表，然后重新将它们转换回numpy数组。这样似乎至少不会失败。这仍然是这种情况吗？ - Derek_M

是的，这样应该可以工作。请注意，tensorflow_set_weights会将所有值转换为numpy数组。 - Maxim

太棒了，我会回复一切的。感谢你帮助我到达那里！ - Derek_M

Answer 4

感谢Maxim帮我找到了解决方案。我想发布一个答案，将图形和权重都转换为JSON，以便那些遇到这个问题的人使用。如果只需要序列化图形而不是权重，我创建了一个要点，概括了Maxim在这里写的内容: Tensorflow graph with non json serialized weights。

现在，为了序列化/反序列化图形和权重，我在这里创建了一个单独的要点: Tensorflow graph with json serialized weights and graph。

为了解释一下，我首先略微调整了权重函数，不返回get weights中的变量，并在set weights中获取当前变量。这是一个重要的注意事项，特别是如果图形与当前可训练变量稍有不同：

import tensorflow as tf
import numpy as np
from google.protobuf import json_format
import json

def tensorflow_get_weights():
    vs = tf.trainable_variables()
    values = tf.get_default_session().run(vs)
    return values

def tensorflow_set_weights(weights):
    assign_ops = []
    feed_dict = {}
    vs = tf.trainable_variables()
    zipped_values = zip(vs, weights)
    for var, value in zipped_values:
        value = np.asarray(value)
        assign_placeholder = tf.placeholder(var.dtype, shape=value.shape)
        assign_op = var.assign(assign_placeholder)
        assign_ops.append(assign_op)
        feed_dict[assign_placeholder] = value
    tf.get_default_session().run(assign_ops, feed_dict=feed_dict)

接下来，我创建了两个实用函数，用于将权重转换为JSON格式和从JSON格式中恢复权重：

def convert_weights_to_json(weights):
    weights = [w.tolist() for w in weights]
    weights_list = json.dumps(weights)
    return weights_list

def convert_json_to_weights(json_weights):
    loaded_weights = json.loads(json_weights)
    loaded_weights = [np.asarray(x) for x in loaded_weights]
    return loaded_weights

然后我有一个方法，最初用于启动训练。这个方法会初始化变量，运行优化器，获取权重和图形，并将它们转换为json格式。代码如下：

def run_initial_with_json_weights(opti, feed_dict):
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for i in range(0, 250):
            sess.run(opti, feed_dict=feed_dict)
        first_weights = tensorflow_get_weights()
        g = tf.get_default_graph().as_graph_def()
    json_string = json_format.MessageToJson(g)
    return json_string, convert_weights_to_json(first_weights)

现在我们有了序列化的权重和图形，如果我们想要继续训练或进行预测，可以按照以下步骤操作。该方法反序列化graphdef和权重，运行优化，然后进行预测。

def run_serialized(json_graph, json_weights, feed_dict):
    gd = tf.GraphDef()
    gd = json_format.Parse(json_graph, gd)
    weights = convert_json_to_weights(json_weights)

    with tf.Session() as sess:
        tf.import_graph_def(gd)
        sess.run(tf.global_variables_initializer())
        nu_out = tf.get_default_graph().get_tensor_by_name('outer/Sigmoid:0')
        mini = tf.get_default_graph().get_tensor_by_name('mini:0')
        tensorflow_set_weights(weights)

        for i in range(0, 50):
             sess.run(mini, feed_dict=feed_dict)

        predicted = sess.run(nu_out, feed_dict=feed_dict)

    return predicted

一个完整的异或（XOR）示例在上面的代码中。

- Derek_M