JS 的 AI 时代来了!

人工智能
有了 JS-Torch​ 之后,在 Node.js、Deno 等 JS Runtime 上跑 AI 应用的日子越来越近了。当然,JS-Torch 要推广起来,它还需要解决一个很重要的问题,即 GPU 加速。

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

责任编辑:武晓燕 来源: 全栈修仙之路
相关推荐

2024-04-12 14:12:45

人工智能JavaScript

2024-04-12 08:32:03

JSAI机器学习

2016-09-02 08:20:33

OpsDevWWDCDevOps

2024-01-29 13:09:39

AI摄影师

2018-03-06 14:59:13

2015-03-30 20:56:55

SaaS云计算企业级应用

2020-06-04 07:00:00

人工智能AIRedis

2012-07-20 15:17:19

移动购物

2015-12-24 10:25:30

微软程序员时代

2021-10-10 19:36:14

AI人工智能

2017-11-16 15:36:02

人工智能云端云计算

2021-05-13 11:35:54

K8STerraform代码仓库

2023-10-09 10:45:35

2021-02-23 09:55:00

AIAI+医疗智能诊断

2019-09-05 19:05:16

2017-12-26 17:22:14

高速公路移动支付收费

2024-03-27 11:12:13

2022-09-26 12:17:56

SIMeSIMiPhone

2023-03-15 08:03:31

点赞
收藏

51CTO技术栈公众号