<reponame>ifiaposto/Tensorflow-Implementation-of-NVAE<filename>util/trainer.py<gh_stars>0
# ============================================================================
#  Copyright 2021. 
#
#
#  Author: <NAME> 
#  Contact: <EMAIL>, <EMAIL> 
#
#
# All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================



"""Tensorflow implementation of NVAE"""

import tensorflow as tf
from tensorflow.keras.callbacks import CSVLogger

import os
import sys
import csv

from util.learning_rate_schedule import LearningRateScheduler,LEARNING_RATE_SCHEDULER
from util.data import DatasetSequence
from util.eval import create_eval_command
from util.hparams import HParams
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True

__all__ = [
           'keras_train',
           "get_default_train_hparams"
           ]

def keras_train(model,
                dataset_builder,
                hparams,
                checkpoint_period,
                train_log_dir,
                train_log_subdir=None,
                epoch=0,
                mirrored_strategy=None,
                ):
    
    """
        utility function that trains a deep variational autoencoder.
        Args:
            model:                  A VAE model to be trained.
            dataset_builder:        A tfds.builder (refer to tensorflow_datasets for the detailed documentation) object.
            hparams:                The training parameters.
            checkpoint_period:      Save period of the model during training.
            train_log_dir:          Directory for the training logs.
            train_log_subdir:       Sub-directory name in train_log_dir for the training logs.
            epoch:                  Epoch from which the training is resumed (if a model_filepath is provided).
            mirrored_strategy:      A tf.distribute.MirroredStrategy in case of a multi-gpu model.
        Output:
            final_model.h5:         Saved model in .5 format. in the specified sub-directory.
            keras_trainer_log.csv:  The learning curves of the deep vae that include: [total loss (including regularization), negative log-likelihood, kl divergence, nelbo]
            eval_command.sh:        Command for loading and evaluation the model.
        
    """
    
    # load dataset
    dataset_builder.download_and_prepare()
    dataset = dataset_builder.as_dataset()    
    train_dataset, test_dataset = dataset["train"], dataset["test"]
    
    assert isinstance(train_dataset, tf.data.Dataset)
    
        
    # extract training samples
    train_dataset = train_dataset.batch(dataset_builder.info.splits['train'].num_examples)
    train_features = tf.compat.v1.data.make_one_shot_iterator(train_dataset).get_next()
    train_images = train_features['image']    
    train_images=tf.dtypes.cast(train_images, tf.float32)
    
    
    # extract test samples
    test_dataset = test_dataset.batch(dataset_builder.info.splits['test'].num_examples)
    test_features = tf.compat.v1.data.make_one_shot_iterator(test_dataset).get_next()
        
    test_images = test_features['image']
    test_images=tf.dtypes.cast(test_images, tf.float32)
        

    # create training directory
    train_log_dir=create_train_log_dir(train_log_dir,hparams, train_log_subdir=model.name if not train_log_subdir else  train_log_subdir)


    ## set-up training checkpoint
    checkpoint=SeReModelCheckpoint(save_freq=checkpoint_period,save_dir=train_log_dir)


   
    training_algorithm[hparams.optimization_method](model=model,
                                                    x=train_images,
                                                    x_test=test_images,
                                                    hparams=hparams,
                                                    checkpoint=checkpoint,
                                                    mirrored_strategy=mirrored_strategy,
                                                    initial_epoch=epoch,
                                                    train_log_dir=train_log_dir,
                                                    )

    
    
    ##### training finished #####
    print('training finished')

    #save model weights
    model.save_weights(train_log_dir+"final_model.h5")



################################################################################################################################################
        ################################################ Optimization Methods ################################################
################################################################################################################################################





def train_by_gradient_descent(model,
                              x,
                              x_test,
                              hparams,
                              checkpoint,
                              mirrored_strategy=None,
                              initial_epoch=0,
                              train_log_dir=None):
    
        """ 
        """
    

        ##### create callbacks  #####



        ## set-up learning rate scheduler, if any

        lr_callback=[LearningRateScheduler(LEARNING_RATE_SCHEDULER[hparams.learning_rate_schedule](initial_learning_rate=hparams.learning_rate,**{key:getattr(hparams, 'lr_'+key) for key in LEARNING_RATE_SCHEDULER[hparams.learning_rate_schedule].__init__.__code__.co_varnames[2:]}))] if hparams.learning_rate_schedule!='constant' else []
    
        #    ## set-up kl- annealing scheduler, if any
        kl_scheduler=KLRegularizerSchedule(hparams.kl_scheduler)
    
        ## set-up inference scheduler (to account for the contribution of the residual decoders in the reconstruction loss), if any
        


        
        ##create loss function
        loss=[lambda y1, y2: y2]*(4)
        
        loss_weight=[ 1.0]+[kl_scheduler.beta]+[0.0,0.0]

        
        # compile model
        print('optimizer')
        print(hparams.optimizer)
        kwargs = {key:getattr(hparams, 'keras_lr_'+key) for key in OPTIMIZER[hparams.optimizer].__init__.__code__.co_varnames[2:-1]}
        



        optimizer=OPTIMIZER[hparams.optimizer](learning_rate=hparams.learning_rate,**kwargs)


        
        if mirrored_strategy is not None:
            with mirrored_strategy.scope():
                print('compile model mirror strategy')
                model.compile(optimizer=optimizer,
                              loss=loss,
                              loss_weights=loss_weight,
                              )

        else:
            print('no mirror strategy')
            model.compile(optimizer=optimizer,
                          loss=loss,
                          loss_weights=loss_weight,
                          )
    
        # start training
        y=[tf.zeros(x.shape[0])]*(4)
        y_test=[tf.zeros(x_test.shape[0])]*(4)
        


        
        model.fit(x=DatasetSequence(x_set=[x], y_set=y, batch_size=hparams.batch_size*hparams.num_gpus, shape=model.image_shape,binary=hparams.binary,flip_horizontally=hparams.flip_horizontally, flip_vertically=hparams.flip_vertically),
                  validation_data=(x_test,y_test),
                  epochs=initial_epoch+hparams.epochs,
                  initial_epoch=initial_epoch,
                  callbacks=[CSVLogger(train_log_dir+'keras_trainer_log.csv',append=True)]+[checkpoint,lr_callback,kl_scheduler],
                  )





OPTIMIZER={"adam":tf.keras.optimizers.Adam,
           "adamax":tf.keras.optimizers.Adamax,
           "rmsprop":tf.keras.optimizers.RMSprop,
}

                      


training_algorithm = {
                      'gradient_descent': train_by_gradient_descent,
                      }

################################################################################################################################################
        ################################################ Utility functions ################################################
################################################################################################################################################

def get_default_train_hparams():
    return HParams(use_mirrored_strategy=False,             # flag whether a multi-gpu model (use_mirrored_strategy=True) will be trained.
                   num_gpus=1,                              # number of available gpus.
                   batch_size=256,                          # batch_size per gpu.
                   epochs=2560,                             # number of training epochs.
                   optimization_method='gradient_descent',  # optimization method (currently the only one supported).
                   kl_scheduler= KLRegularizerSchedule.get_default_hparams(), # parameters for the kl-annealing.
                   learning_rate=1e-2,                      # initial learning rate for the rest of the model
                   
                   ## parameters for the learning rate scheduler
                   learning_rate_schedule='cosine_warmup',       # learning rate scheduler
                   lr_name='learning_rate_scheduler',            # name of the scheduler.
                   lr_decay_epochs=1000,                         # number of epochs, after the warmup, during which the learning rate decreases from learning_rate to lr_min_learning_rate at the end of each epoch
                   lr_steps_per_epoch=1563,                      # the learning steps per gpu, applied during a epoch.
                   lr_verbose=False,                             # flag indicating whether the current learning rate will be printed.
                   lr_warmup_epochs=5,                           # number of epochs, during which the learning rate increases from lr_min_learning_rate to learning_rate per gradient step.
                   lr_min_learning_rate=0.0001,                  # minimum learning rate.
                   ## parameters for the optimizer.
                   ## refer to the documentation of tf.keras.optimizers.Adam for a detailed description.
                   optimizer="adam",
                   keras_lr_momentum=0.0001,
                   keras_lr_beta_1=0.9,                           # the exponential decay rate for the 1st moment estimates.
                   keras_lr_beta_2=0.999,                         # the exponential decay rate for the 2nd moment estimates. 
                   keras_lr_epsilon=1e-03,                        # small constant for numerical stability.
                   keras_lr_kwargs='',
                   keras_lr_name='',
                   ## options for the dataset augmentation per training step
                   binary=False,
                   flip_horizontally=True,
                   flip_vertically=True,
                   )


def create_train_log_dir(train_log_dir, hparams, train_log_subdir=None):
    """ utility function that creates a unique subdirectory in train_log_dir with postfix name train_log_subdir.
    """
    
    if train_log_subdir:
        i = 0
        while os.path.exists(train_log_dir+str(i)+"_"+train_log_subdir+"/"):
            i += 1

        train_log_dir=train_log_dir+str(i)+"_"+train_log_subdir+"/"
        os.makedirs(train_log_dir, exist_ok=True)

    # save the command line arguments
    f= open(train_log_dir+"train.sh","w")
    f.write('python3 '+' '.join(sys.argv[0:]))
    f.close()


    # save the training hyperparameters
    with open(train_log_dir+"train_hparams.csv", 'w', newline='') as csv_file:
        csv_hparams_logger = csv.writer(csv_file, delimiter=',',quotechar='|', quoting=csv.QUOTE_MINIMAL)
        csv_hparams_logger.writerow(['training hyperparams'])
        hparams.save(csv_hparams_logger)
    
    
  
    
    # create the command for the model evaluation
    create_eval_command(sys.argv[0:],train_log_dir)
    

    return train_log_dir



class SeReModelCheckpoint(tf.keras.callbacks.Callback):
    """
        utility class for periodically saving the model.
    """
    
    def __init__(self, save_freq, save_dir):
    
        self.save_freq = save_freq
    
        self.save_dir=save_dir
    
    def on_epoch_end(self, epoch, logs=None):
        print('finished epoch: %d, learning rate: %s'%(epoch,tf.keras.backend.eval(self.model.optimizer.lr)))

    def on_epoch_end(self, epoch, logs=None):
        
            if (epoch+1) % self.save_freq == 0: #and self.model.nof_active_layers==self.model.nof_layers:
                
                # create and save command for resuming the training
                f= open(self.save_dir+"restore_train.sh","w")
                f.write('python3 '+' '.join(sys.argv[0:])+' --epoch='+str(epoch+1)+' --model_filepath='+self.save_dir+'/checkpoint_model.h5')
                f.close()
                cloned_model=self.model.init()
                cloned_model.save_weights("./"+self.save_dir+"checkpoint_model.h5")



class KLRegularizerSchedule(tf.keras.callbacks.Callback):
    
    """
        utility function for kl-annealing.
        
        Returns the kl-regularization term  for the current epoch needed for the deterministic warm-up described in the paper:
        [1]. <NAME>., <NAME>., <NAME>., <NAME>. and <NAME>., 2016. Ladder variational autoencoders.
        In Advances in neural information processing systems (pp. 3738-3746).
        [2]. <NAME>., <NAME>., <NAME>., <NAME>. and <NAME>., 2016.
            How to train deep variational autoencoders and probabilistic ladder networks.
            In 33rd International Conference on Machine Learning (ICML 2016),
        [3]. <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>.
             Generating sentences from a continuous space. arXiv preprint arXiv:1511.06349. 2015 Nov 19
        
    """
    def __init__(self, hparams):
        self.beta = tf.Variable(0.0)

        self.hparams = hparams

    def on_epoch_begin(self, epoch, logs=None):
   

        if self.hparams.update_beta=='smooth':
            if epoch<self.hparams.warmup_epochs:
                if epoch % self.hparams.warmup_smooth_period !=0:
                    beta= tf.keras.backend.eval(model.beta)
                beta= epoch*((1.0-self.hparams.warmup_smooth_start)/self.hparams.warmup_epochs)+self.hparams.warmup_smooth_start
            else:
                beta= 1.0
        elif self.hparams.update_beta=='hard':
            if epoch<self.hparams.warmup_epochs:
                beta=  0.0
            else:
                beta= 1.0
        tf.keras.backend.update(self.beta,beta)

    @staticmethod
    def get_default_hparams():
        return HParams(
                       warmup_epochs=0,         # number of epochs for the deterministic warm-up (KL regularization term <1), see [2].
                       update_beta='smooth',    # scheme for the update of the KL-regularization term.
                       warmup_smooth_start=0.0, # offset value for linear warm-up (start value of the kl-coefficient for epoch=0).
                       warmup_smooth_period=1,  # number of epochs very which the kl coefficient is updated.
                       )



# ============================================================================
#  Copyright 2021. 
#
#
#  Author: <NAME> 
#  Contact: <EMAIL>, <EMAIL> 
#
#
# All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================



"""Tensorflow implementation of NVAE"""

import tensorflow as tf
from tensorflow.keras.callbacks import CSVLogger

import os
import sys
import csv

from util.learning_rate_schedule import LearningRateScheduler,LEARNING_RATE_SCHEDULER
from util.data import DatasetSequence
from util.eval import create_eval_command
from util.hparams import HParams
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True

__all__ = [
           'keras_train',
           "get_default_train_hparams"
           ]

def keras_train(model,
                dataset_builder,
                hparams,
                checkpoint_period,
                train_log_dir,
                train_log_subdir=None,
                epoch=0,
                mirrored_strategy=None,
                ):
    
    """
        utility function that trains a deep variational autoencoder.
        Args:
            model:                  A VAE model to be trained.
            dataset_builder:        A tfds.builder (refer to tensorflow_datasets for the detailed documentation) object.
            hparams:                The training parameters.
            checkpoint_period:      Save period of the model during training.
            train_log_dir:          Directory for the training logs.
            train_log_subdir:       Sub-directory name in train_log_dir for the training logs.
            epoch:                  Epoch from which the training is resumed (if a model_filepath is provided).
            mirrored_strategy:      A tf.distribute.MirroredStrategy in case of a multi-gpu model.
        Output:
            final_model.h5:         Saved model in .5 format. in the specified sub-directory.
            keras_trainer_log.csv:  The learning curves of the deep vae that include: [total loss (including regularization), negative log-likelihood, kl divergence, nelbo]
            eval_command.sh:        Command for loading and evaluation the model.
        
    """
    
    # load dataset
    dataset_builder.download_and_prepare()
    dataset = dataset_builder.as_dataset()    
    train_dataset, test_dataset = dataset["train"], dataset["test"]
    
    assert isinstance(train_dataset, tf.data.Dataset)
    
        
    # extract training samples
    train_dataset = train_dataset.batch(dataset_builder.info.splits['train'].num_examples)
    train_features = tf.compat.v1.data.make_one_shot_iterator(train_dataset).get_next()
    train_images = train_features['image']    
    train_images=tf.dtypes.cast(train_images, tf.float32)
    
    
    # extract test samples
    test_dataset = test_dataset.batch(dataset_builder.info.splits['test'].num_examples)
    test_features = tf.compat.v1.data.make_one_shot_iterator(test_dataset).get_next()
        
    test_images = test_features['image']
    test_images=tf.dtypes.cast(test_images, tf.float32)
        

    # create training directory
    train_log_dir=create_train_log_dir(train_log_dir,hparams, train_log_subdir=model.name if not train_log_subdir else  train_log_subdir)


    ## set-up training checkpoint
    checkpoint=SeReModelCheckpoint(save_freq=checkpoint_period,save_dir=train_log_dir)


   
    training_algorithm[hparams.optimization_method](model=model,
                                                    x=train_images,
                                                    x_test=test_images,
                                                    hparams=hparams,
                                                    checkpoint=checkpoint,
                                                    mirrored_strategy=mirrored_strategy,
                                                    initial_epoch=epoch,
                                                    train_log_dir=train_log_dir,
                                                    )

    
    
    ##### training finished #####
    print('training finished')

    #save model weights
    model.save_weights(train_log_dir+"final_model.h5")



################################################################################################################################################
        ################################################ Optimization Methods ################################################
################################################################################################################################################





def train_by_gradient_descent(model,
                              x,
                              x_test,
                              hparams,
                              checkpoint,
                              mirrored_strategy=None,
                              initial_epoch=0,
                              train_log_dir=None):
    
        """ 
        """
    

        ##### create callbacks  #####



        ## set-up learning rate scheduler, if any

        lr_callback=[LearningRateScheduler(LEARNING_RATE_SCHEDULER[hparams.learning_rate_schedule](initial_learning_rate=hparams.learning_rate,**{key:getattr(hparams, 'lr_'+key) for key in LEARNING_RATE_SCHEDULER[hparams.learning_rate_schedule].__init__.__code__.co_varnames[2:]}))] if hparams.learning_rate_schedule!='constant' else []
    
        #    ## set-up kl- annealing scheduler, if any
        kl_scheduler=KLRegularizerSchedule(hparams.kl_scheduler)
    
        ## set-up inference scheduler (to account for the contribution of the residual decoders in the reconstruction loss), if any
        


        
        ##create loss function
        loss=[lambda y1, y2: y2]*(4)
        
        loss_weight=[ 1.0]+[kl_scheduler.beta]+[0.0,0.0]

        
        # compile model
        print('optimizer')
        print(hparams.optimizer)
        kwargs = {key:getattr(hparams, 'keras_lr_'+key) for key in OPTIMIZER[hparams.optimizer].__init__.__code__.co_varnames[2:-1]}
        



        optimizer=OPTIMIZER[hparams.optimizer](learning_rate=hparams.learning_rate,**kwargs)


        
        if mirrored_strategy is not None:
            with mirrored_strategy.scope():
                print('compile model mirror strategy')
                model.compile(optimizer=optimizer,
                              loss=loss,
                              loss_weights=loss_weight,
                              )

        else:
            print('no mirror strategy')
            model.compile(optimizer=optimizer,
                          loss=loss,
                          loss_weights=loss_weight,
                          )
    
        # start training
        y=[tf.zeros(x.shape[0])]*(4)
        y_test=[tf.zeros(x_test.shape[0])]*(4)
        


        
        model.fit(x=DatasetSequence(x_set=[x], y_set=y, batch_size=hparams.batch_size*hparams.num_gpus, shape=model.image_shape,binary=hparams.binary,flip_horizontally=hparams.flip_horizontally, flip_vertically=hparams.flip_vertically),
                  validation_data=(x_test,y_test),
                  epochs=initial_epoch+hparams.epochs,
                  initial_epoch=initial_epoch,
                  callbacks=[CSVLogger(train_log_dir+'keras_trainer_log.csv',append=True)]+[checkpoint,lr_callback,kl_scheduler],
                  )





OPTIMIZER={"adam":tf.keras.optimizers.Adam,
           "adamax":tf.keras.optimizers.Adamax,
           "rmsprop":tf.keras.optimizers.RMSprop,
}

                      


training_algorithm = {
                      'gradient_descent': train_by_gradient_descent,
                      }

################################################################################################################################################
        ################################################ Utility functions ################################################
################################################################################################################################################

def get_default_train_hparams():
    return HParams(use_mirrored_strategy=False,             # flag whether a multi-gpu model (use_mirrored_strategy=True) will be trained.
                   num_gpus=1,                              # number of available gpus.
                   batch_size=256,                          # batch_size per gpu.
                   epochs=2560,                             # number of training epochs.
                   optimization_method='gradient_descent',  # optimization method (currently the only one supported).
                   kl_scheduler= KLRegularizerSchedule.get_default_hparams(), # parameters for the kl-annealing.
                   learning_rate=1e-2,                      # initial learning rate for the rest of the model
                   
                   ## parameters for the learning rate scheduler
                   learning_rate_schedule='cosine_warmup',       # learning rate scheduler
                   lr_name='learning_rate_scheduler',            # name of the scheduler.
                   lr_decay_epochs=1000,                         # number of epochs, after the warmup, during which the learning rate decreases from learning_rate to lr_min_learning_rate at the end of each epoch
                   lr_steps_per_epoch=1563,                      # the learning steps per gpu, applied during a epoch.
                   lr_verbose=False,                             # flag indicating whether the current learning rate will be printed.
                   lr_warmup_epochs=5,                           # number of epochs, during which the learning rate increases from lr_min_learning_rate to learning_rate per gradient step.
                   lr_min_learning_rate=0.0001,                  # minimum learning rate.
                   ## parameters for the optimizer.
                   ## refer to the documentation of tf.keras.optimizers.Adam for a detailed description.
                   optimizer="adam",
                   keras_lr_momentum=0.0001,
                   keras_lr_beta_1=0.9,                           # the exponential decay rate for the 1st moment estimates.
                   keras_lr_beta_2=0.999,                         # the exponential decay rate for the 2nd moment estimates. 
                   keras_lr_epsilon=1e-03,                        # small constant for numerical stability.
                   keras_lr_kwargs='',
                   keras_lr_name='',
                   ## options for the dataset augmentation per training step
                   binary=False,
                   flip_horizontally=True,
                   flip_vertically=True,
                   )


def create_train_log_dir(train_log_dir, hparams, train_log_subdir=None):
    """ utility function that creates a unique subdirectory in train_log_dir with postfix name train_log_subdir.
    """
    
    if train_log_subdir:
        i = 0
        while os.path.exists(train_log_dir+str(i)+"_"+train_log_subdir+"/"):
            i += 1

        train_log_dir=train_log_dir+str(i)+"_"+train_log_subdir+"/"
        os.makedirs(train_log_dir, exist_ok=True)

    # save the command line arguments
    f= open(train_log_dir+"train.sh","w")
    f.write('python3 '+' '.join(sys.argv[0:]))
    f.close()


    # save the training hyperparameters
    with open(train_log_dir+"train_hparams.csv", 'w', newline='') as csv_file:
        csv_hparams_logger = csv.writer(csv_file, delimiter=',',quotechar='|', quoting=csv.QUOTE_MINIMAL)
        csv_hparams_logger.writerow(['training hyperparams'])
        hparams.save(csv_hparams_logger)
    
    
  
    
    # create the command for the model evaluation
    create_eval_command(sys.argv[0:],train_log_dir)
    

    return train_log_dir



class SeReModelCheckpoint(tf.keras.callbacks.Callback):
    """
        utility class for periodically saving the model.
    """
    
    def __init__(self, save_freq, save_dir):
    
        self.save_freq = save_freq
    
        self.save_dir=save_dir
    
    def on_epoch_end(self, epoch, logs=None):
        print('finished epoch: %d, learning rate: %s'%(epoch,tf.keras.backend.eval(self.model.optimizer.lr)))

    def on_epoch_end(self, epoch, logs=None):
        
            if (epoch+1) % self.save_freq == 0: #and self.model.nof_active_layers==self.model.nof_layers:
                
                # create and save command for resuming the training
                f= open(self.save_dir+"restore_train.sh","w")
                f.write('python3 '+' '.join(sys.argv[0:])+' --epoch='+str(epoch+1)+' --model_filepath='+self.save_dir+'/checkpoint_model.h5')
                f.close()
                cloned_model=self.model.init()
                cloned_model.save_weights("./"+self.save_dir+"checkpoint_model.h5")



class KLRegularizerSchedule(tf.keras.callbacks.Callback):
    
    """
        utility function for kl-annealing.
        
        Returns the kl-regularization term  for the current epoch needed for the deterministic warm-up described in the paper:
        [1]. <NAME>., <NAME>., <NAME>., <NAME>. and <NAME>., 2016. Ladder variational autoencoders.
        In Advances in neural information processing systems (pp. 3738-3746).
        [2]. <NAME>., <NAME>., <NAME>., <NAME>. and <NAME>., 2016.
            How to train deep variational autoencoders and probabilistic ladder networks.
            In 33rd International Conference on Machine Learning (ICML 2016),
        [3]. <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>.
             Generating sentences from a continuous space. arXiv preprint arXiv:1511.06349. 2015 Nov 19
        
    """
    def __init__(self, hparams):
        self.beta = tf.Variable(0.0)

        self.hparams = hparams

    def on_epoch_begin(self, epoch, logs=None):
   

        if self.hparams.update_beta=='smooth':
            if epoch<self.hparams.warmup_epochs:
                if epoch % self.hparams.warmup_smooth_period !=0:
                    beta= tf.keras.backend.eval(model.beta)
                beta= epoch*((1.0-self.hparams.warmup_smooth_start)/self.hparams.warmup_epochs)+self.hparams.warmup_smooth_start
            else:
                beta= 1.0
        elif self.hparams.update_beta=='hard':
            if epoch<self.hparams.warmup_epochs:
                beta=  0.0
            else:
                beta= 1.0
        tf.keras.backend.update(self.beta,beta)

    @staticmethod
    def get_default_hparams():
        return HParams(
                       warmup_epochs=0,         # number of epochs for the deterministic warm-up (KL regularization term <1), see [2].
                       update_beta='smooth',    # scheme for the update of the KL-regularization term.
                       warmup_smooth_start=0.0, # offset value for linear warm-up (start value of the kl-coefficient for epoch=0).
                       warmup_smooth_period=1,  # number of epochs very which the kl coefficient is updated.
                       )
# ============================================================================ #  Copyright 2021. # # #  Author: <NAME> #  Contact: <EMAIL>, <EMAIL> # # # All Rights Reserved. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ Tensorflow implementation of NVAE utility function that trains a deep variational autoencoder.         Args:             model:                  A VAE model to be trained.             dataset_builder:        A tfds.builder (refer to tensorflow_datasets for the detailed documentation) object.             hparams:                The training parameters.             checkpoint_period:      Save period of the model during training.             train_log_dir:          Directory for the training logs.             train_log_subdir:       Sub-directory name in train_log_dir for the training logs.             epoch:                  Epoch from which the training is resumed (if a model_filepath is provided).             mirrored_strategy:      A tf.distribute.MirroredStrategy in case of a multi-gpu model.         Output:             final_model.h5:         Saved model in .5 format. in the specified sub-directory.             keras_trainer_log.csv:  The learning curves of the deep vae that include: [total loss (including regularization), negative log-likelihood, kl divergence, nelbo]             eval_command.sh:        Command for loading and evaluation the model. # load dataset # extract training samples # extract test samples # create training directory ## set-up training checkpoint ##### training finished ##### #save model weights ################################################################################################################################################ ################################################ Optimization Methods ################################################ ################################################################################################################################################ ##### create callbacks  ##### ## set-up learning rate scheduler, if any #    ## set-up kl- annealing scheduler, if any ## set-up inference scheduler (to account for the contribution of the residual decoders in the reconstruction loss), if any ##create loss function # compile model # start training ################################################################################################################################################ ################################################ Utility functions ################################################ ################################################################################################################################################ # flag whether a multi-gpu model (use_mirrored_strategy=True) will be trained. # number of available gpus. # batch_size per gpu. # number of training epochs. # optimization method (currently the only one supported). # parameters for the kl-annealing. # initial learning rate for the rest of the model ## parameters for the learning rate scheduler # learning rate scheduler # name of the scheduler. # number of epochs, after the warmup, during which the learning rate decreases from learning_rate to lr_min_learning_rate at the end of each epoch # the learning steps per gpu, applied during a epoch. # flag indicating whether the current learning rate will be printed. # number of epochs, during which the learning rate increases from lr_min_learning_rate to learning_rate per gradient step. # minimum learning rate. ## parameters for the optimizer. ## refer to the documentation of tf.keras.optimizers.Adam for a detailed description. # the exponential decay rate for the 1st moment estimates. # the exponential decay rate for the 2nd moment estimates. # small constant for numerical stability. ## options for the dataset augmentation per training step utility function that creates a unique subdirectory in train_log_dir with postfix name train_log_subdir. # save the command line arguments # save the training hyperparameters # create the command for the model evaluation utility class for periodically saving the model. #and self.model.nof_active_layers==self.model.nof_layers: # create and save command for resuming the training utility function for kl-annealing.                  Returns the kl-regularization term  for the current epoch needed for the deterministic warm-up described in the paper:         [1]. <NAME>., <NAME>., <NAME>., <NAME>. and <NAME>., 2016. Ladder variational autoencoders.         In Advances in neural information processing systems (pp. 3738-3746).         [2]. <NAME>., <NAME>., <NAME>., <NAME>. and <NAME>., 2016.             How to train deep variational autoencoders and probabilistic ladder networks.             In 33rd International Conference on Machine Learning (ICML 2016),         [3]. <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>.              Generating sentences from a continuous space. arXiv preprint arXiv:1511.06349. 2015 Nov 19 # number of epochs for the deterministic warm-up (KL regularization term <1), see [2]. # scheme for the update of the KL-regularization term. # offset value for linear warm-up (start value of the kl-coefficient for epoch=0). # number of epochs very which the kl coefficient is updated.