Js-pytorch NPM

PyTorch in JavaScript

JS-PyTorch is a Deep Learning JavaScript library built from scratch, to closely follow PyTorch's syntax.
This library has GPU support, using GPU.js.
If you want to run it yourself, check out the Documentation.
Try out the Web Demo!

Note: You can install the package locally with: npm install js-pytorch

1.Table of Contents

2. Installation

On MacOS, Windows, and Ubuntu, you can install the library with npm install js-pytorch.
On Windows, if you run into an error, you might need to install the latest version of Visual Studio, including the "Desktop development with C++" workload.
To run in the Browser, npm install the latest version of JS-PyTorch, and link the distribution files in your HTML file:

<script src="./node_modules/js-pytorch/dist/utils.js"></script>
<script type="module">
  import { torch } from './node_modules/js-pytorch/dist/js-pytorch-browser.js';
  window.torch = torch;
</script>

After that, you can use JS-PyTorch freely in any <script> in your HTML file.

3. Running it Yourself

Simple Autograd Example:

const { torch } = require("js-pytorch");

// Pass device as an argument to a Tensor or nn.Module (same as PyTorch):
const device = 'gpu';

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

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

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

// Get gradients from specific Tensors:
console.log(w.grad);
console.log(b.grad);

Complex Autograd Example (Transformer):

const { torch } = require("js-pytorch");
const nn = torch.nn;
const optim = torch.optim;
const device = 'gpu';

// Define training hyperparameters:
const vocab_size = 52;
const hidden_size = 32;
const n_timesteps = 16;
const n_heads = 4;
const dropout_p = 0;
const batch_size = 8;

// Create Transformer decoder Module:
class Transformer extends nn.Module {
  constructor(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p, device) {
    super();
    // Instantiate Transformer's Layers:
    this.embed = new nn.Embedding(vocab_size, hidden_size);
    this.pos_embed = new nn.PositionalEmbedding(n_timesteps, hidden_size);
    this.b1 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_p, device);
    this.b2 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_p, device);
    this.ln = new nn.LayerNorm(hidden_size);
    this.linear = new nn.Linear(hidden_size, vocab_size, device);
  }

  forward(x) {
    let z;
    z = torch.add(this.embed.forward(x), this.pos_embed.forward(x));
    z = this.b1.forward(z);
    z = this.b2.forward(z);
    z = this.ln.forward(z);
    z = this.linear.forward(z);
    return z;
  }
}

// Instantiate your custom nn.Module:
const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p, device);

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

// Instantiate sample input and output:
let x = torch.randint(0, vocab_size, [batch_size, n_timesteps, 1]);
let y = torch.randint(0, vocab_size, [batch_size, n_timesteps]);
let loss;

// Training Loop:
for (let i = 0; i < 40; i++) {
  // Forward pass through the Transformer:
  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 loss at every iteration:
  console.log(`Iter ${i} - Loss ${loss.data[0].toFixed(4)}`)
}

Saving and Loading models:

// Instantiate your model:
const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p);

// Train the model:
trainModel(model);

// Save model to JSON file:
torch.save(model, 'model.json')

// To load, instantiate placeHolder using the original model's architecture:
const placeHolder = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p);

// Load weights into placeHolder:
const newModel = torch.load(placeHolder, 'model.json')

4. Distribution & Devtools

Build for Distribution by running npm run build. CJS and ESM modules and index.d.ts will be output in the dist/ folder.
Check the Code with ESLint at any time, running npm run lint.
Run tests run npm test.
Improve Code Formatting with prettier, running npm run prettier.
Performance Benchmarks are also included in the tests/benchmarks/ directory. Run all benchmarks with npm run bench and save new benchmarks with npm run bench:update.

5. Future Work

This package is not as optimized as PyTorch yet, but I tried making it more interpretable. Efficiency improvements are incoming!
Feel free to contribute! Create a merge request to the develop branch, and also feel free to reach out. I'll try to answer as soon as possible.
Hope you enjoy!