How to create a simple TensorFlow.js project

November 25, 2023November 25, 2023 Alian Yu AI

This instruction shows how to create a simple project to predict the y value of equation y = 1.2x + 5 based on TensorFLow.js

Create a webpage

Create a webpage ‘predict.html’:

<html>
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>Traffic Flow Predictor</title>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
</head>
<body>
    <div id="output"></div>
    <script>
        // Create Training Data
        const xs = tf.tensor([0, 1, 2, 3, 4]);
        const ys = xs.mul(1.2).add(5);

        // Output xs and ys
        document.getElementById('output').innerText = xs.dataSync();
        document.getElementById('output').innerText+= "\n" + ys.dataSync();
    </script>
</body>
</html>

Use web browser to open ‘predict.html’, if the TensorFlow.js is loaded correctly, you should see the following text:

0,1,2,3,4
5,6.199999809265137,7.400000095367432,8.600000381469727,9.800000190734863

Train the model

Use the following code to train the model.

    // Define a Linear Regression Model
    const model = tf.sequential();
    model.add(tf.layers.dense({ units: 1, inputShape: [1] }));

    // Specify Loss and Optimizer
    model.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });

    // Train the Model
    model.fit(xs, ys, { epochs: 500 }).then(() => { trainingCompleted() });

    // Training completed
    function trainingCompleted() {
        document.getElementById('output').innerText += "\n" + "Model training completed!";
    }

Use the model to predict the value of y when x = 10

// Use the model
function predict() {
    let result = model.predict(tf.tensor([Number(10)]));
    result.data().then(y => {
        document.getElementById('output').innerText += "\n" + Number(y);
    });
}

Source Code

<html>
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>Traffic Flow Predictor</title>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
</head>
<body>
    <div id="output"></div>
    <script>
        // Create Training Data
        const xs = tf.tensor([0, 1, 2, 3, 4]);
        const ys = xs.mul(1.2).add(5);

        // Output xs and ys
        document.getElementById('output').innerText = xs.dataSync();
        document.getElementById('output').innerText += "\n" + ys.dataSync();
        
        // Define a Linear Regression Model
        const model = tf.sequential();
        model.add(tf.layers.dense({ units: 1, inputShape: [1] }));

        // Specify Loss and Optimizer
        model.compile({ loss: 'meanSquaredError', optimizer: 'sgd' });

        // Train the Model
        model.fit(xs, ys, { epochs: 500 }).then(() => { trainingCompleted() });

        // Training completed
        function trainingCompleted() {
            document.getElementById('output').innerText += "\n" + "Model training completed!";

            predict();
        }

        // Use the model
        function predict() {
            let result = model.predict(tf.tensor([Number(10)]));
            result.data().then(y => {
                document.getElementById('output').innerText += "\n" + Number(y);
            });
        }
    </script>
</body>
</html>

Leave a Reply Cancel reply