Tensorflow.js: Simple linear regression not working great

Question

I have this basic code from an example video (first half of the video does this but with a different data set).

Code:



    
        
        
        
        Website
        
        
        
    
    
        Welcome to my website.

I train it with 100 values where the input is the same as the output. Then when I try running it with 50 as input after training, I get results that range from like -50 to 60.

Is this a normal thing? I would expect values close to 50.

Also, when I train with an array with values from 1 to 1000, and input 500, I even get ouputs starting at -600.

user_5627599 · Answer

There are two major problems in your code.

Problem 1: Not waiting for the Promise to finish

linearModel.fit returns a Promise that will be resolved when training is finished. That means, right now, the model will start training, but before the training is finished you are already asking for predictions.

You have to wait for the Promise to resolve. Easiest way to do that is to put your code into an async function and use await like this:

(async () => {
  var linearModel = tf.sequential();
  // ...
  await linearModel.fit(xs, ys);
  // ...
})();

Problem 2: Having a bad learning rate

The default learning rate for sgd is 0.01, which is too high in your case. Training your model with that value increased the loss value for each training iteration for me indicating that the learning rate is to high. If you are want more information on the topic check out this guide on "estimating an optimal learning rate".

You can change the learning rate by using the tf.train.sgd function (instead of using a string) and passing a learning rate like this:

linearModel.compile({ loss: 'meanSquaredError', optimizer: tf.train.sgd(0.0001) });

Further improvement: Train for more than one epoch

Although, the above tips should already yield results that are close to 50, you can further improve your model by training for more than one epoch. You can pass the epochs parameter like this to train your model for longer:

await linearModel.fit(xs, ys, {
  epochs: 10
});

Try it yourself

I added the improvements to your code in the below snippet. You can change the epochs and learning rate values to see how it affects the resulting prediction of 50.

document.querySelector('button').addEventListener('click', async () => {
    const learningRate = document.querySelector('#learning_rate').value;
    const epochs = document.querySelector('#epochs').value;
    
    const linearModel = tf.sequential();
    linearModel.add(tf.layers.dense({ units: 1, inputShape: [1] }));
    linearModel.compile({ loss: 'meanSquaredError', optimizer: tf.train.sgd(learningRate) });

    const xs = tf.tensor1d([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100]);
    const ys = tf.tensor1d([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100]);

    await linearModel.fit(xs, ys, {
        epochs,
        callbacks: {
            onEpochEnd: (epoch, logs) => console.log(`Loss, epoch ${epoch}: ${logs.loss}`),
        },
    });

    function linearPrediction(val) {
        const output = linearModel.predict(tf.tensor1d([val]));
        const prediction = Array.from(output.dataSync())[0];
        console.log(`Prediction for 50: ${prediction}`);
    }

    linearPrediction(50);
});



epochs: 
learning rate:

Tensorflow.js: Simple linear regression not working great

1 Answers

Problem 1: Not waiting for the Promise to finish

Problem 2: Having a bad learning rate

Further improvement: Train for more than one epoch

Try it yourself