Search code examples
tensorflowhyperparameters

How to set hyperparameter (learning_rate) schedule in TensorFlow?


What is the way to schedule hyper-parameters in TensorFlow?

Namely, for the sake of reproducibility I would like to implement a ResNet (you name one) using suggested learning rate schedule {0: 0.1, 1: 1., 100: 0.01, 150: 0.001}, or enable the weight decay only after first few initial epoch.

For example, tensorpack provides an optionas follows:

ScheduledHyperParamSetter('learning_rate', [(1, 0.1), (82, 0.01), (123, 0.001), (300, 0.0002)])

How can that be done in native TF?


Solution

  • Ok, it wasn't that hard

        schedule = {1: 0.1, 2: 0.2, 3: 0.3, 4: 0.4, 100: 0.01, 150: 0.001}
        schedule = sorted(config.lr_schedule.items(), key=lambda x: x[0])
    
        boundaries = [num_train_iter * int(x[0]) for x in schedule]
        rates = [x[1] for x in schedule]
        rates = rates[:1] + rates  # 
        assert len(boundaries) + 1 == len(rates)
    
        learning_rate = tf.train.piecewise_constant(tf.cast(global_step, tf.int32), boundaries, rates)