tensorflow - 在 TensorFlow BYOL 架构中加入指数移动平均线

问题描述

我正在尝试在 TensorFlow 中重新创建BYOL 架构以执行表示学习任务。考虑以下来自Tian 等人的示意图。(2021) :

在左侧，您会看到一个名为“Online”的模型组件，其权重为 W。在右侧，您会看到一个名为“Target”的模型组件，其权重为 Wa。目标组件实际上是在线组件的指数移动平均 (EMA)，因此权重 Wa 不会被任何基于梯度的优化器更新（因此是停止梯度）。当我查看 and 的文档时tf.train.ExponentialMovingAverage，tfa.optimizers.MovingAverage我发现这些示例主要集中在将 EMA 应用于一个模型组件（例如本示例中的在线组件）而不是其影子（目标组件）。

我正在努力理解如何：

...将目标组件实例化为在线组件的副本
... 对给定输入执行在线和目标组件的评估
...将目标组件的权重更新为在线组件的指数移动平均值

我在想上述可以通过存储trainable_variables在线组件的组合来实现，然后将目标组件的那些分配给在线组件，执行推理，然后换回旧的trainable_variables，并在两者之间执行一些更新。或者，我可以两次实例化在线组件，并trainable_variables使用梯度下降更新第一个组件，并使用 EMA 更新第二个组件。但我不知道如何将其转换为工作代码。任何帮助将不胜感激。

标签： tensorflowkeras

tensorflow - 在 TensorFlow BYOL 架构中加入指数移动平均线

问题描述

解决方案

推荐阅读