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

Question

我想对模型的 2 个部分使用预训练的权重。我有 2 个来自不同模型的检查点，当我使用估算器架构时，我只能使用 tf.estimator.WarmStart 将一个检查点加载到我的主模型中。

tf.WarmStartSettings(ckpt_to_initialize_from=X)

from the doc：

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

我看不出如何添加额外的检查点。也许有一种方法可以将两个检查点的权重加载到一个并加载那个？

Answer 1

您可以使用init_from_checkpoint。

首先，定义assignment map：

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 中。它将覆盖标准变量初始化。