如何在张量流上更新随机权重平均（SWA）中的权重？

Question

我对如何实现 tfa 的 SWA optimizer 感到困惑。这里有两点：

1. 当您查看文档时，它会将您指向 [this] 模型平均教程。该教程使用 tfa.callbacks.AverageModelCheckpoint，它允许您

• 将移动平均权重分配给模型并保存。
• （或）保留旧的非平均权重，但保存的模型使用平均权重。

拥有一个独特的 ModelCheckpoint 可以让您保存移动平均权重（而不是当前权重）是有意义的。但是 - 似乎 SWA 应该管理权重平均。这让我想设置update_weights=False。

这是正确的吗？本教程使用update_weights=True。

1. documentation 中有一条关于 SWA 未更新 BN 层的说明。按照here 的建议，我这样做了，

# original training
model.fit(...)
# updating weights from final run 
optimizer.assign_average_vars(model.variables)
# batch-norm-hack: lr=0 as suggested https://stackoverflow.com/a/64376062/607528
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0),
    loss=loss,
    metrics=metrics)
model.fit(
    data,
    validation_data=None,
    epochs=1,
    callbacks=final_callbacks)

在保存我的模型之前。

这对吗？

谢谢！

Answer 1

处理批范数的最简单方法如下： 首先，遍历模型中的所有层并重置批范数层中的移动均值和移动方差（在我的示例中，我假设批范数层以“bn”结尾）：

for l in model.layers:
    if l.name.split('_')[-1] == 'bn': # e.g. conv1_bn
       l.moving_mean.assign(tf.zeros_like(l.moving_mean))
       l.moving_variance.assign(tf.ones_like(l.moving_variance))

之后，运行您的模型一个时期并将训练设置为 true 以更新移动平均值和方差：

count = 0
for x,_ in dataset_train:
    _ = model(x, training = True)
    count += 1
    if count > steps_per_epoch:
       break

Answer 2

有两种方法可以做到这一点，第一种是您在保存之前手动更新权重，例如文档中的example。

import tensorflow as tf
import tensorflow_addons as tfa

model = tf.Sequential([...])
opt = tfa.optimizers.SWA(
        tf.keras.optimizers.SGD(lr=2.0), 100, 10)
model.compile(opt, ...)
model.fit(x, y, ...)

# Update the weights to their mean before saving
opt.assign_average_vars(model.variables)

model.save('model.h5')

如果您设置了update_weights = True，则第二个选项是通过 AverageModelCheckpoint 更新权重。正如协作笔记本example 所示

avg_callback = tfa.callbacks.AverageModelCheckpoint(filepath=checkpoint_dir, 
                                                    update_weights=True)
...

#Build Model
model = create_model(moving_avg_sgd)

#Train the network
model.fit(fmnist_train_ds, epochs=5, callbacks=[avg_callback])

注意AverageModelCheckpoint 在保存模型之前也会从source code 调用assign_average_vars：

def _save_model(self, epoch, logs):
    optimizer = self._get_optimizer()
    assert isinstance(optimizer, AveragedOptimizerWrapper)

    if self.update_weights:
        optimizer.assign_average_vars(self.model.variables)
        return super()._save_model(epoch, logs)
    ...