使用tf.estimator.WarmStartSettings加载（或合并）几个预先训练的检查点

Question

我想对模型的2个部分使用预训练权重。我有2个来自不同模型的检查点，使用tf.estimator.WarmStart只能从其中一个装载到我的主模型中，因为我正在使用estimator架构。

tf.WarmStartSettings(ckpt_to_initialize_from=X)

from the doc：

  

可以提供目录或特定的检查点（对于前者，将使用最新的检查点）。

我看不到如何添加其他检查点。也许有一种方法可以将两个检查点的权重都加载到一个，然后再加载一个？

Answer 1

您可以使用init_from_checkpoint。

首先，定义分配图：

dir = 'path_to_checkpoint_files'
vars_to_load = [i[0] for i in tf.train.list_variables(dir)]

这将创建检查点中所有变量的列表

assignment_map = {variable.op.name: variable for variable in tf.global_variables() if variable.op.name in vars_to_load}

这将创建一个字典，将当前图形中的变量作为键，并将检查点中的变量作为值

tf.train.init_from_checkpoint(dir, assignment_map)

此函数位于估算器的model_fn中。它将覆盖标准变量的初始化。