使用 tf.contrib.graph_editor 克隆网络

Question

我有一个像这样的函数来构建网络。

def build_network(inputs):
  # Some arbitrary set of variables and ops here. For example...
  out = tf.contrib.layers.fully_connected(inputs, 123)
  (...)
  return out

然后我用它来构建这样的网络。

inputs = tf.placeholder(...)
outputs = build_network(inputs)

如果我想构建更多具有相同结构但独立变量的网络，我只需要在其他变量范围和可选的其他输入下再次调用 build_network。

我的问题是：如果这个 build_network 不再可用，但原始网络的输入和输出可用，我该怎么办？换句话说：如何将整个子图从 outputs 一直克隆到 inputs 到另一个具有自己独立变量集但结构相同的变量范围中？

我的理解是 tf.contrib.graph_editor 和特别是 graph_editor.copy 正是我需要做这些事情的工具。但是，我找不到任何使用它们的好例子。有什么建议吗？

Answer 1

注意：这个答案和OP's answer 是相辅相成的。请先阅读OP's answer。

我今天在这个问题上花了 4 个小时。这是 TensorFlow 丑陋的地方之一（这就是为什么你应该使用 PyTorch 如果你喜欢图形操作）。

---

这里的关键点是：tf.Variable 不是图形元素（更多关于它的here），而是围绕 3 个操作的包装器：Assign 操作，Read op 和VariableV2 op 本质上是ref tensor（更多关于它here）。因此，您需要在 TensorFlow 框架中显式调用它。

如果我们仔细查看graph_editor 的代码，尤其是transform module，我们可以看到它只在tf.Graph 上运行，而不涉及TensorFlow 框架中的任何内容。因此，graph_editor.copy（和类似的）方法根本不会触及tf.Variable 对象。它只复制构成tf.Variable 构建块的张量和操作。

好的，那我们如何解决这个问题呢？

假设你有以下变量：

var = tf.get_trainable_variables()[0]
print(var.to_proto())
# variable_name: "dense_1/kernel:0"
# initializer_name: "dense_1/kernel/Assign"
# snapshot_name: "dense_1/kernel/read:0"
# initial_value_name: "dense_1/random_uniform:0"
# trainable: true

您知道在graph_editor.copy(...) 之后，您的dense_1 名称范围现在是dense_1b。然后，您只需使用info.transformed(...) 获取相应的操作和张量，然后执行以下操作：

from tensorflow.core.framework import variable_pb2

var_def = variable_pb2.VariableDef()
var_def.variable_name = 'dense_1b/kernel:0'
var_def.initializer_name = "dense_1b/kernel/Assign"
var_def.snapshot_name = "dense_1b/kernel/read:0"
var_def.initial_value_name = "dense_1/random_uniform:0"
var_def.trainable = True

现在，我想强调tf.Variable documentation的以下部分：

variable_def: ... 使用其内容重新创建变量对象，引用图中必须已经存在的变量节点。图表没有改变。

因此，tf.Variable 构造函数允许我们在现有的图形元素之上创建一个变量包装器。这正是我们所需要的：

cloned_var = tf.Variable(variable_def=var_def)
tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, cloned_var)

解决了！

---

我使这个答案尽可能简单和具体，以展示tf.Variables 的基本机制。您现在可以轻松实现更一般情况下的代码以自动生成新变量。

PS：我讨厌 TensorFlow！

Answer 2

回应自己，我发现了一种复制子图的方法。

from tensorflow.contrib import graph_editor as ge

# From the example above.
inputs = [tf.placeholder(...), ...]
outputs = build_network(inputs)

sgv = ge.make_view(ge.get_within_boundary_ops(
    tf.get_default_graph(),
    [t.op for t in outputs],
    [t.op for t in inputs]))

# This could be any new inputs. In this example I build new identical placeholders.
new_inputs = {p: tf.placeholder(dtype=p.dtype, shape=p.shape) for p in inputs}
new_sgv, info = ge.copy_with_input_replacements(sgv, new_inputs, dst_scope='copy')

new_inputs = [info.transformed(t) for t in inputs]
new_outputs = [info.transformed(t) for t in outputs]

但是，现在我在尝试使用网络副本时遇到了一个新问题。副本中的新变量未初始化，尝试运行 tf.global_variables_initializer() 也无济于事。

原因是因为它们的 tf.Variable 从未构建过，所以它们不是 GlobalKeys.GLOBAL_VARIABLES 集合的一部分。我可以轻松找到与这些变量对应的操作以及它们在原始和副本之间的映射，但我无法从中构建 tf.Variable。

我发现了一些用于初始化的 hacky 解决方法，但它仅适用于集合中的 var。

init_ops = []
for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
  if v.op in sgv.ops:
    init_ops.append(info.transformed(v.initializer))

...

session.run([tf.global_variables_initializer()] + init_ops)

有没有更好的方法来做到这一点？理想情况下，允许为复制的变量创建 tf.Variables 以将它们添加到全局变量集合中。或者，如果这不可能，至少是一种可靠的方法来获取初始化操作，而无需找到原始网络的 tf.Variable 对象。