BatchRenormCallback.py

BatchRenormCallback.py

class BatchRenormCallback(tf.keras.callbacks.Callback):
    def __init__(self, log_dir, bvalues):
        super().__init__()
        self.writer = tf.summary.create_file_writer(log_dir)
        self.bvalues = bvalues

    def on_epoch_begin(self, epoch, logs=None):
        if epoch in self.bvalues:
            self.change_renorm(epoch)

    def change_renorm(self, epoch):
        print(f"Changing renorm clipping values epoch {epoch}", self.bvalues[epoch])
        renorm_clipping = self.bvalues[epoch]

        with self.writer.as_default():
            tf.summary.scalar("renorm_rmax", renorm_clipping["rmax"], epoch)
            tf.summary.scalar("renorm_dmax", renorm_clipping["dmax"], epoch)

        for layer in self.model.layers:
            if isinstance(layer, tf.keras.Model):
                for layer2 in layer.layers:
                    if isinstance(layer2, BatchNormalization):
                        layer2.renorm_clipping = renorm_clipping
            if isinstance(layer, BatchNormalization):
                layer.renorm_clipping = renorm_clipping