How do you set the Adam optimizer learning rate in tensorflow.js?

Question

For tensorflow.js, how do you set the learning rate for the Adam optimizer in node.js? I get an error:

model.optimizer.setLearningRate is not a function

const optimizer = tf.train.adam(0.001)
model.compile({
    loss: 'sparseCategoricalCrossentropy',
    optimizer,
    shuffle: true,
    metrics: ['accuracy']
});

await model.fit(trainValues, trainLabels, {
    epochs: 50,
    validationData: [testValues, testLabels],
    callbacks: {
        onEpochBegin: async (epoch) => {
            const newRate = getNewRate();
            model.optimizer.setLearningRate(newRate);
        }
    }
});

Answer 1

When you call model.compile, you can pass an instance of tf.train.Optimizer instead of passing a string. These instances are created via the tf.train.* factories and you can pass the learning rate as first argument.

Code Sample

model.compile({
    optimizer: tf.train.sgd(0.000001), // custom learning rate
    /* ... */
});

Changing the Learning Rate during training

Currently, only the sgd optimizers has a setLearningRate method implemented, meaning the following code only works for optimizer instances created via tf.train.sgd:

const optimizer = tf.train.sgd(0.001);
optimizer.setLearningRate(0.000001);

Using non-official API

The optimizer instances are having a protected attribute learningRate, which you can change. The attribute is not public but, as this is JavaScript, you can simply change the value by setting learningRate on the object like this:

const optimizer = tf.train.adam();
optimizer.learningRate = 0.000001;
// or via your model:
model.optimizer.learningRate = 0.000001;

Keep in mind, you are using a non-official part of the API, which might break anytime.

Answer 2

When creating a model, one can set the learning rate when passing the optimizer to model.compile

const myOptimizer = tf.train.sgd(myLearningRate) 
model.compile({optimizer: myOptimizer, loss: 'meanSquaredError'});

The learning rate can be changed during training by using setLearningRate

model.fit(xs, ys, {
  epochs: 800, 
  callbacks: {
    onEpochEnd: async (epoch, logs) => {

      if (epoch == 300){
        model.optimizer.setLearningRate(0.14)  
    }

      if (epoch == 400){
        model.optimizer.setLearningRate(0.02)   
      }
    }
  } 
})