随着人工智能技术的不断进步，javascript 也迎来了自己的 ai 时代。
js-torch 是一个全新的深度学习库，专为 javascript 设计，其语法习惯与广受欢迎的 pytorch 框架高度相似。这个库提供了一整套深度学习工具，包括可追踪梯度的张量对象、多层网络结构以及自动微分功能。
pytorch，作为 js-torch 的灵感来源，是一个由 meta ai 团队开发的开源深度学习框架。它以简洁、灵活和易用著称，其动态计算图让神经网络的构建变得更加直观。
js-torch 可以通过 npm 或 pnpm 进行安装，也可以在线体验其提供的 demo。
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 等 javascript 运行时环境中运行 ai 应用铺平了道路。
