Deep Learning in Javascript
Tags: machinelearning, javascript, ai, showdev
Published: 2024-04-01 14:55:24

- [JS-PyTorch](https://github.com/eduardoleao052/js-pytorch) is a Deep Learning __JavaScript library__ built from scratch, to closely follow __PyTorch__'s syntax. 
- It is implemented from scratch in a __well-documented, unit tested, and interpretable__ way,  so it can help other JavaScript learners get into Machine Learning!

- Feel free to try out a <a href="https://eduardoleao052.github.io/js-pytorch/assets/demo/demo.html" target="blank">Web Demo</a>!

<p> &nbsp; </p>

## 1. TL;DR

-  `src/tensor.ts` contains a fully functional __Tensor__ object, which can track gradients.
- `src/layers.ts` contains many Deep Learning __Layers__ and functions.

> **Note:** The [project's README](https://github.com/eduardoleao052/js-pytorch) contains details on all available operations and layers. 

<p> &nbsp; </p>

## 2. Running it Yourself
### Install & Import
To start off, you can install the package locally running `npm install js-pytorch` on the terminal.
Then, on your JavaScript file, import the package with:
```javascript
const torch = require("js-pytorch");
```

<p> &nbsp; </p>

### Create Tensors
To use all of these cool deep learning __Tensor Operations__, we need to instantiate some Tensors:

```javascript
// Instantiate Tensors:
let x = torch.randn([8,4,5]);
let w = torch.randn([8,5,4], requires_grad = true);
let b = torch.tensor([0.2, 0.5, 0.1, 0.0], requires_grad = true);
```

The syntax is the same as __PyTorch__'s: 
- `torch.tensor(Array)` recieves an array and turns it into a __Tensor__.
- `torch.randn(shape)` creates a Tensor filled with normally-distributed random numbers, with the provided shape.
- The `requires_grad` argument is set to `true` if we want to __optimize__ this parameter (by tracking it's gradients).

<p> &nbsp; </p>

### Tensor Operations
Now, let's run some operations on these Tensors:

```javascript
// Make calculations:
let out = torch.matmul(x, w);
out = torch.add(out, b);
```

- `torch.matmul(x, w)` performs matrix multiplication between `x` and `w` (just like in __PyTorch__).
- `torch.add(out, b)` adds both Tensors.

> **Note:** As `w` has `require_grad` set to `true`, its __children__ `out` will also have it's gradients tracked.

<p> &nbsp; </p>

### Getting Gradients

```javascript
// Compute gradients on whole graph:
out.backward();

// Get gradients from specific Tensors:
console.log(w.grad);
console.log(b.grad);
```
 
- Calling `out.backward()` calculates the gradients of every Tensor that let to it (its __parents__), relative to `out`.
- IRL, we will call `.backward()` on the `loss` Tensor, to get the gradients necessary to reduce it.
- To access a Tensor's gradients, simply call `Tensor.grad`.

<p> &nbsp; </p>

## 3. Full Example (Neural Network): 
In this example, we implement a full __Neural Network__, with three Linear layers, and ReLU activations. The syntax for the `nn.Module` class is identical to [PyTorch](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)'s.
```javascript
const torch = require("js-pytorch");
const nn = torch.nn;
const optim = torch.optim;

// Implement Module class:
class NeuralNet extends nn.Module {
  constructor(in_size, hidden_size, out_size) {
    super();
    // Instantiate Neural Network's Layers:
    this.w1 = new nn.Linear(in_size, hidden_size);
    this.relu1 = new nn.ReLU();
    this.w2 = new nn.Linear(hidden_size, hidden_size);
    this.relu2 = new nn.ReLU();
    this.w3 = new nn.Linear(hidden_size, out_size);
  };

  forward(x) {
    let z;
    z = this.w1.forward(x);
    z = this.relu1.forward(z);
    z = this.w2.forward(z);
    z = this.relu2.forward(z);
    z = this.w3.forward(z);
    return z;
  };
};

// Instantiate Model:
let in_size = 16;
let hidden_size = 32;
let out_size = 10;
let batch_size = 16;

let model = new NeuralNet(in_size,hidden_size,out_size);

// Define loss function and optimizer:
let loss_func = new nn.CrossEntropyLoss();
let optimizer = new optim.Adam(model.parameters(), 3e-3);

// Instantiate input and output:
let x = torch.randn([batch_size, in_size]);
let y = torch.randint(0, out_size, [batch_size]);
let loss;

// Training Loop:
for (let i = 0; i < 256; i++) {
  let z = model.forward(x);

  // Get loss:
  loss = loss_func.forward(z, y);

  // Backpropagate the loss using torch.tensor's backward() method:
  loss.backward();

  // Update the weights:
  optimizer.step();

  // Reset the gradients to zero after each training step:
  optimizer.zero_grad();

  // Print current loss:
  console.log(`Iter: ${i} - Loss: ${loss.data}`);
}
```

<p> &nbsp; </p>

## 4. Building for distribution & DevTools

- To __build__ for distribution, run `npm run build`. CJS and ESM modules and `index.d.ts` will be output in the `dist/` folder.
- To check the code with __ESLint__ at any time, run`npm run lint`.
- To improve code formatting with __prettier__, run `npm run prettier`.

## 5. Conclusion
- Hope you enjoyed the package! 
- New additions, such as __GPU support__ are coming soon.