Как динамически обновлять норму партии в TF2?

Я нашел реализацию PyTorch, которая уменьшает параметр пакетной нормы momentum с 0.1 в первой эпохе до 0.001 в последней эпохе. Любые предложения о том, как сделать это с параметром нормы партии momentum в TF2? (т. е. начинать с 0.9 и заканчивать на 0.999). Например, вот что делается в коде PyTorch:

# in training script
momentum = initial_momentum * np.exp(-epoch/args.epochs * np.log(initial_momentum/final_momentum))
model_pos_train.set_bn_momentum(momentum)

# model class function
def set_bn_momentum(self, momentum):
    self.expand_bn.momentum = momentum
    for bn in self.layers_bn:
        bn.momentum = momentum

РЕШЕНИЕ:

Выбранный ниже ответ обеспечивает жизнеспособное решение при использовании tf.keras.Model.fit() API. Тем не менее, я использовал пользовательский цикл обучения. Вот что я сделал вместо этого:

После каждой эпохи:

mi = 1 - initial_momentum  # i.e., inital_momentum = 0.9, mi = 0.1
mf = 1 - final_momentum  # i.e., final_momentum = 0.999, mf = 0.001
momentum = 1 - mi * np.exp(-epoch / epochs * np.log(mi / mf))
model = set_bn_momentum(model, momentum)

Функция set_bn_momentum (спасибо этой статье):

def set_bn_momentum(model, momentum):
    for layer in model.layers:
        if hasattr(layer, 'momentum'):
            print(layer.name, layer.momentum)
            setattr(layer, 'momentum', momentum)

    # When we change the layers attributes, the change only happens in the model config file
    model_json = model.to_json()

    # Save the weights before reloading the model.
    tmp_weights_path = os.path.join(tempfile.gettempdir(), 'tmp_weights.h5')
    model.save_weights(tmp_weights_path)

    # load the model from the config
    model = tf.keras.models.model_from_json(model_json)

    # Reload the model weights
    model.load_weights(tmp_weights_path, by_name=True)
    return model

Этот метод не добавлял значительных накладных расходов в программу обучения.