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

• Installation
• Running it Yourself
  • Simple Autograd Example
  • Complex Autograd Example (Transformer)
  • Saving and Loading models
• Distribution & Devtools
• Future Work

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="https://github.com/eduardoleao052/js-pytorch/raw/main/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!