python – 将特定TensorFlow变量还原到特定层(按名称还原)

2023年2月5日 229次阅读

假设我训练了TensorFlow模型并保存它,现在有一个不同的模型,我想在模型中的某些层中使用保存模型中的一些权重(它们具有相同的形状).

现在,我能够找到如何从模型中保存特定变量(具有特定名称),但我无法找到任何按名称恢复这些变量的示例.

例如,假设在我保存的模型中,我保存了称为“v1”的张量(具有某种形状).现在在我的新模型中,我有一个称为“v2”的权重张量(它具有与“v1”张量相同的形状).现在我想将保存的变量“v1”加载到我的“v2”权重张量,或者甚至更好地将这个“v1”权重加载到我的新图中的多个张量.

这甚至可能吗？如果是这样,我该怎么办？

最佳答案 我找到了解决此问题的解决方法.

您可以做的是将变量值直接保存到磁盘,可以是值,也可以是带有密钥作为张量名称和值的字典.
例如：

vars_dict = {}
for tensor in (list_of_tensors_you_want_to_save):
     vars_dict[tensor.name] = sess.run(tensor)

然后,您可以将此字典中的任何变量加载到所需的任何其他变量.

在我的例子中,假设原始张量被称为“v1”,而我要加载的两个张量是“v2”和“v3”,可以完成以下操作：

tensor_to_load_1 = tf.get_default_graph().get_tensor_by_name("v2")
tensor_to_load_2 = tf.get_default_graph().get_tensor_by_name("v3")

assign_op_1 = tf.assign(tensor_to_load_1, vars_dict["v1"])
assign_op_1 = tf.assign(tensor_to_load_2, vars_dict["v1"])

sess.run([assign_op_1, assign_op_2])

当然,这仅受“v1”,“v2”和“v3”必须具有相同形状的事实的限制.

使用此示例代码,您可以保存所有变量并将其加载到您想要的任何其他变量,而无需使用原始图表来匹配您当前的变量.

点赞