用 JavaScript
做深度学习 Deep Learning in
JavaScript

// 用约 10 行代码训练神经网络// Train a neural network in ~10 lines
const torch = require('estorch');

const model = new SimpleNet();
const criterion = new torch.nn.CrossEntropyLoss();
const optimizer = new torch.optim.Adam(model.parameters(), { lr: 0.001 });

for await (const [images, labels] of trainLoader) {
  optimizer.zeroGrad();
  const outputs = model.forward(images);
  const loss = criterion.forward(outputs, labels);
  loss.backward();
  optimizer.step();
}

// 创建支持自动微分的张量并进行运算// Create and manipulate tensors with autograd
const a = torch.tensor([[1, 2], [3, 4]], { requiresGrad: true });
const b = torch.tensor([[5, 6], [7, 8]], { requiresGrad: true });

// 矩阵乘法// Matrix multiplication
const c = a.matmul(b);

// 逐元素运算// Element-wise operations
const d = c.relu().mul(0.5);

// 反向传播// Backpropagation
d.backward();
console.log(a.grad);  // 梯度自动计算// gradients automatically computed

// 通过继承 nn.Module 定义模型// Define a model by extending nn.Module
class SimpleNet extends torch.nn.Module {
  constructor() {
    super();
    this.conv1 = new torch.nn.Conv2d(1, 8, 3, { padding: 1 });
    this.pool  = new torch.nn.AvgPool2d(2);
    this.fc1   = new torch.nn.Linear(14 * 14 * 8, 64);
    this.fc2   = new torch.nn.Linear(64, 10);
  }

  forward(x) {
    x = torch.nn.functional.relu(this.conv1.forward(x));
    x = this.pool.forward(x);
    x = torch.flatten(x, 1);
    x = torch.nn.functional.relu(this.fc1.forward(x));
    return this.fc2.forward(x);
  }
}