npm install js-pytorch
pnpm add js-pytorch

目前，JS-Torch 已经支持了基本的张量操作，如加法、减法、乘法和除法等，以及一些常用的深度学习层，例如nn.Linear、nn.MultiHeadSelfAttention、nn.FullyConnected、nn.Block等。

以下是一个简单的 JS-Torch 使用示例，展示了如何进行自动梯度计算：

// 导入 torch 模块
import { torch } from "js-pytorch";

// 创建张量
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 });

// 执行计算
let out = torch.matmul(x, w);
out = torch.add(out, b);

// 计算梯度
out.backward();

// 输出梯度
console.log(w.grad);
console.log(b.grad);

更复杂的使用示例，如 Transformer 模型的实现，也包含在 JS-Torch 中：

// 导入 torch 模块和 nn 模块
import { torch, nn } from "js-pytorch";

class Transformer extends nn.Module {
  constructor(vocab_size, hidden_size, n_timesteps, n_heads, p) {
    super();
    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 = 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;
  }
}

// 创建模型实例
const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p);

// 定义损失函数和优化器
const loss_func = new nn.CrossEntropyLoss();
const optimizer = new optim.Adam(model.parameters(), { lr: 5e-3, reg: 0 });

// 创建样本输入和输出
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;
// 训练循环
for (let i = 0; i < 40; i++) {
  // 通过 Transformer 模型进行前向传播
  let z = model.forward(x);
  // 计算损失
  loss = loss_func.forward(z, y);
  // 使用 torch.tensor 的 backward 方法反向传播损失
  loss.backward();
  // 更新权重
  optimizer.step();
  // 每个训练步骤后将梯度重置为零
  optimizer.zero_grad();
}

JS-Torch 为在 Node.js、Deno * 运行时环境中运行 AI 应用铺平了道路。

本文转载自：博客园

原文：