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