如何在tensorflow上的随机加权平均（SWA）中更新权重？

Question

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

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

• 将移动平均权重指定给模型，并保存它们
• （或）保留旧的非平均权重，但保存的模型使用平均权重

有一个独特的ModelCheckpoint，允许您保存移动平均权重（而不是当前权重）是有意义的。然而，似乎SWA应该管理平均体重。这使我想设置update_weights=False

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

1. 有一个关于SWA不更新documentation中的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

有两种方法可以做到这一点，第一种方法是在保存之前手动更新权重，如文档中的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')

第二个选项是，如果设置了^{，则通过AverageModelCheckpoint更新权重。如collab笔记本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)
    ...

Answer 2

处理批次定额的最简单方法如下： 首先，循环模型中的所有层，并重置批次标准层中的移动平均值和移动方差（在我的示例中，我假设批次标准层以“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))

之后，运行模型一个历元并将training设置为true以更新移动平均值和方差：

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