Блог

Я пытаюсь воссоздать архитектуру BYOL в TensorFlow для задачи обучения представлению. Рассмотрим следующую принципиальную схему из Tian et al. (2021):

Слева вы видите компонент модели под названием «Онлайн» с весами W. Справа вы видите компонент модели под названием «Цель» с весами Wa. Целевой компонент на самом деле является экспоненциальной скользящей средней (EMA) онлайн-компонента, поэтому веса Wa не обновляются никаким оптимизатором на основе градиента (отсюда и стоп-градиент). Когда я просматриваю документацию tf.train.ExponentialMovingAverage и tfa.optimizers.MovingAverage , я вижу, что примеры в значительной степени сосредоточены на применении EMA к одному компоненту модели (например, онлайн-компоненту в этом примере), а не к его тени (целевой компонент).

Я изо всех сил пытаюсь понять, как:

1. … создайте экземпляр целевого компонента как копию сетевого компонента
2. … выполните оценку онлайн-и целевых компонентов для данного ввода
3. … обновите веса целевого компонента в виде экспоненциальной скользящей средней онлайн-компонента

Я думал, что вышесказанное может быть достигнуто за счет сочетания сохранения trainable_variables онлайн-компонента, затем назначения компонентов целевого компонента онлайн-компоненту, выполнения вывода , затем замены старого trainable_variables и выполнения некоторых обновлений между ними. В качестве альтернативы я мог бы создать экземпляр онлайн-компонента дважды и обновить trainable_variables первый с помощью градиентного спуска, а второй-с помощью EMA. Но я бы не знал, как преобразовать это в рабочий код. Любая помощь была бы очень признательна.