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

Question

我想对模型的两个部分使用预训练权重。我有两个来自不同模型的检查点，从中我只能用tf.estimator.WarmStart公司因为我使用的是估计器架构。在

tf.WarmStartSettings(ckpt_to_initialize_from=X)

from the doc：

Either the directory or a specific checkpoint can be provided (in the case of the former, the latest checkpoint will be used).

我不知道如何添加额外的检查点。也许有一种方法可以将权重从两个检查点加载到一个检查点，然后再加载那个检查点？在

Answer 1

您可以使用init_from_checkpoint。在

首先，定义任务映射：

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

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

^{pr2}$

这将创建一个dict，它将当前图中的变量作为键，将检查点中的变量作为值

tf.train.init_from_checkpoint(dir, assignment_map)

这个函数放在估计器的model_fn内。它将覆盖标准变量初始化。在