Home > Article > Web Front-end > Getting Started with Machine Learning in JavaScript: A Beginner’s Guide with TensorFlow.js
Machine learning (ML) has rapidly transformed the world of software development. Until recently, Python was the dominant language in the ML space, thanks to libraries like TensorFlow and PyTorch. But with the rise of TensorFlow.js, JavaScript developers can now dive into the exciting world of machine learning, using familiar syntax to build and train models directly in the browser or on Node.js.
In this blog post, we’ll explore how you can get started with machine learning using JavaScript. We’ll walk through an example of building and training a simple model using TensorFlow.js.
TensorFlow.js is an open-source library that allows you to define, train, and run machine learning models entirely in JavaScript. It runs both in the browser and on Node.js, making it incredibly versatile for a wide range of ML applications.
Here are a few reasons why TensorFlow.js is exciting:
Let’s see how to get started!
Before diving into code, you’ll need to install TensorFlow.js. You can include it in your project via a <script> tag or npm, depending on your environment.</p> <h3> Browser Setup </h3> <p>To use TensorFlow.js in the browser, simply include the following <script> tag in your HTML file:<br> </p> <pre class="brush:php;toolbar:false"><script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script> </pre> <h3> Node.js Setup </h3> <p>For a Node.js environment, you can install it using npm:<br> </p> <pre class="brush:php;toolbar:false">npm install @tensorflow/tfjs </pre> <h2> 2. Building a Simple Neural Network Model </h2> <p>Let’s create a simple neural network that predicts the output of a basic linear function, y = 2x - 1. We will use TensorFlow.js to create and train this model.</p> <h3> Step 1: Define the Model </h3> <p>We’ll start by defining a sequential model (a linear stack of layers) with one dense layer:<br> </p> <pre class="brush:php;toolbar:false">// Import TensorFlow.js import * as tf from '@tensorflow/tfjs'; // Create a simple sequential model const model = tf.sequential(); // Add a single dense layer with 1 unit (neuron) model.add(tf.layers.dense({units: 1, inputShape: [1]})); </pre> <p>Here, we’ve created a model with one dense layer. The layer has one neuron (units: 1), and it expects a single input feature (inputShape: [1]).</p> <h3> Step 2: Compile the Model </h3> <p>Next, we compile the model by specifying the optimizer and loss function:<br> </p> <pre class="brush:php;toolbar:false">// Compile the model model.compile({ optimizer: 'sgd', // Stochastic Gradient Descent loss: 'meanSquaredError' // Loss function for regression }); </pre> <p>We use the Stochastic Gradient Descent (SGD) optimizer, which is effective for small models. The loss function, meanSquaredError, is appropriate for regression tasks like this one.</p> <h3> Step 3: Prepare the Training Data </h3> <p>We’ll now create some training data for the function y = 2x - 1. In TensorFlow.js, data is stored in tensors (multidimensional arrays). Here's how we can generate some training data:<br> </p> <pre class="brush:php;toolbar:false">// Generate some synthetic data for training const xs = tf.tensor2d([0, 1, 2, 3, 4], [5, 1]); // Inputs (x values) const ys = tf.tensor2d([1, 3, 5, 7, 9], [5, 1]); // Outputs (y values) </pre> <p>In this case, we’ve created a tensor xs with input values (0, 1, 2, 3, 4) and a corresponding output tensor ys with values calculated using y = 2x - 1.</p> <h3> Step 4: Train the Model </h3> <p>Now, we can train the model on our data:<br> </p> <pre class="brush:php;toolbar:false">// Train the model model.fit(xs, ys, {epochs: 500}).then(() => { // Once training is complete, use the model to make predictions model.predict(tf.tensor2d([5], [1, 1])).print(); // Output will be close to 2*5 - 1 = 9 }); </pre> <p>Here, we train the model for 500 epochs (iterations over the training data). After training, we use the model to predict the output for an input value of 5, which should return a value close to 9 (y = 2*5 - 1 = 9).</p> <h2> 3. Running the Model in the Browser </h2> <p>To run this model in the browser, you’ll need an HTML file that includes the TensorFlow.js library and your JavaScript code:<br> </p> <pre class="brush:php;toolbar:false"><!DOCTYPE html> <html lang="en"> <head> <meta charset="UTF-8"> <meta name="viewport" content="width=device-width, initial-scale=1.0"> <title>TensorFlow.js Example</title> <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script> <script defer src="app.js"></script>