Tensorflow.js inconsistent prediction, returns 0 or works as intended - tensorflow.js

Question

I'm trying to do a simple Tensorflow.js linear model but I get inconsistent results. It will either return 0 for any input value entered, or it will work as intended (for example if I enter 11 for the input it returns close to 110).
When the page loads up it either works or it doesn't. I can get it to work if I refresh the page 3 or 4 times. Once it works it seems to keep working.
What am I doing wrong?
import {Component, OnInit} from '#angular/core';
import * as tf from '#tensorflow/tfjs';
#Component({
selector: 'app-linear-model',
templateUrl: './linear-model.component.html',
styleUrls: ['./linear-model.component.css']
})
export class LinearModelComponent implements OnInit {
title = 'Linear Model example';
linearModel: tf.Sequential;
prediction: any;
xData: number[] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
yData: number[] = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100];
constructor() {
}
ngOnInit() {
this.trainNewModel();
}
async trainNewModel() {
// this is based on the following tutorial:
// https://angularfirebase.com/lessons/tensorflow-js-quick-start/#Step-2-Install-Tensorflow-js
const learningRate = 0.01;
const optimizerVar = tf.train.sgd(learningRate);
// Define a model for linear regression.
this.linearModel = tf.sequential();
this.linearModel.add(tf.layers.dense({units: 1, inputShape: [1], activation: 'relu'}));
// Prepare the model for training: Specify the loss and the optimizer.
this.linearModel.compile({loss: 'meanSquaredError', optimizer: optimizerVar});
// Training data defined at top
const x = tf.tensor1d(this.xData);
const y = tf.tensor1d(this.yData);
// Train
await this.linearModel.fit(x, y, {epochs: 10});
console.log('model trained!');
}
predict(val) {
val = parseFloat(val);
const output = this.linearModel.predict(tf.tensor2d([val], [1, 1])) as any;
this.prediction = Array.from(output.dataSync())[0];
console.log(output.toString());
}
}

Answer 1

Your issue has to do with the random initialization of the kernel of the dense layer.
Given the value of the weights and the bias, it might happen that the learning rate causes the loss not to decrease. One can keep track of the loss value and if that happens to lower the learning rate.
Another way to solve the issue would be to set an initializer matrix for the dense layers.
this.linearModel.add(tf.layers.dense({units: 1, inputShape: [1], activation: 'relu', kernelInitializer:'ones'}
live code here