JS 的 AI 时代来了！

JS-Torch 简介

JS-Torch[1] 是一个从零开始构建的深度学习 JavaScript 库，其语法与 PyTorch[2] 非常接近。它包含一个功能齐全的张量对象（可跟踪梯度）、深度学习层和函数，以及一个自动微分引擎。

图片

PyTorch 是一个开源的深度学习框架，由 Meta 的人工智能研究团队开发和维护。它提供了丰富的工具和库，用于构建和训练神经网络模型。PyTorch 的设计理念是简单、灵活，以及易于使用，它的动态计算图特性使得模型的构建更加直观和灵活。

你可以通过 npm 或 pnpm 来安装 js-pytorch:

npm install js-pytorch
pnpm add js-pytorch

或者在线体验 js-pytorch 提供的 Demo[3]：

图片

https://eduardoleao052.github.io/js-torch/assets/demo/demo.html

JS-Torch 已支持的功能

目前 JS-Torch 已经支持 Add、Subtract、Multiply、Divide 等张量操作，同时也支持Linear、MultiHeadSelfAttention、ReLU 和 LayerNorm 等常用的深度学习层。

Tensor Operations

Add
Subtract
Multiply
Divide
Matrix Multiply
Power
Square Root
Exponentiate
Log
Sum
Mean
Variance
Transpose
At
MaskedFill
Reshape

Deep Learning Layers

nn.Linear
nn.MultiHeadSelfAttention
nn.FullyConnected
nn.Block
nn.Embedding
nn.PositionalEmbedding
nn.ReLU
nn.Softmax
nn.Dropout
nn.LayerNorm
nn.CrossEntropyLoss

JS-Torch 使用示例

Simple Autograd

import { torch } from "js-pytorch";

// 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));

// 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 (Transformer)

import { torch } from "js-pytorch";
const nn = torch.nn;

class Transformer extends nn.Module {
  constructor(vocab_size, hidden_size, n_timesteps, n_heads, p) {
    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 = p)
    );
    this.b2 = new nn.Block(
      hidden_size,
      hidden_size,
      n_heads,
      n_timesteps,
      (dropout_p = p)
    );
    this.ln = new nn.LayerNorm(hidden_size);
    this.linear = new nn.Linear(hidden_size, vocab_size);
  }

  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
);

// 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();
}

有了 JS-Torch 之后，在 Node.js、Deno 等 JS Runtime 上跑 AI 应用的日子越来越近了。当然，JS-Torch 要推广起来，它还需要解决一个很重要的问题，即 GPU 加速。目前已有相关的讨论，如果你感兴趣的话，可以进一步阅读相关内容：GPU Support[4] 。

参考资料

[1]JS-Torch: https://github.com/eduardoleao052/js-torch

[2]PyTorch: https://pytorch.org/

[3]Demo: https://eduardoleao052.github.io/js-torch/assets/demo/demo.html

[4]GPU Support: https://github.com/eduardoleao052/js-torch/issues/1