Загрузите (или объедините) несколько предварительно обученных контрольных точек с помощью tf.estimator.WarmStartSettings

Question

Я хочу использовать предварительно подготовленные грузы для 2 частей моей модели.У меня есть 2 контрольные точки от разных моделей, из которых я могу загрузить только одну в основную модель с помощью tf.estimator.WarmStart, поскольку я использую архитектуру оценки.

tf.WarmStartSettings(ckpt_to_initialize_from=X)

из документа :

Можно указать либо каталог, либо конкретную контрольную точку (в случае первой будет использоваться последняя контрольная точка).

Я могу 'не вижу, как я могу добавить дополнительную контрольную точку.Может быть, есть способ загрузить веса с обеих контрольных точек в одну и загрузить эту?

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}

Иэто создает dict, который имеет переменные из текущего графа в качестве ключа и переменные из контрольной точки в качестве значений

tf.train.init_from_checkpoint(dir, assignment_map)

Эта функция находится внутри оценки model_fn.Это заменит стандартную переменную инициализации.