All Examples
Example 13
advanced
13
Neural Networks
Autograd
Training

Neural Network Training

This example demonstrates the full neural network training loop in Deepbox. You build a Sequential model (Linear → ReLU → Linear), create training data as GradTensors via parameter(), define an Adam optimizer, and run a training loop: forward pass → compute MSE loss → backward() to get gradients → optimizer.step() to update weights → optimizer.zeroGrad() to reset. You also learn to build custom Module subclasses with explicit parameter registration, switch between train/eval modes, serialize model state with stateDict(), and use SGD with momentum as an alternative optimizer. The example covers 4 parts: autograd training, custom modules, plain tensor inference, and SGD training.

Deepbox Modules Used

deepbox/ndarraydeepbox/nndeepbox/optim

What You Will Learn

  • Build Sequential models from Linear, ReLU, and other layers
  • Use parameter() to create GradTensors that track gradients
  • Training loop: forward → loss → backward → optimizer.step → zeroGrad
  • Use Adam for fast convergence; SGD with momentum for fine-tuning
  • Switch to plain tensor .forward() for inference (no graph overhead)

Source Code

13-neural-network-training/index.ts
1import { GradTensor, parameter, tensor } from "deepbox/ndarray";2import { Linear, Module, mseLoss, ReLU, Sequential } from "deepbox/nn";3import { Adam, SGD } from "deepbox/optim";45console.log("=== Neural Network Training ===\n");67// Build model8const model = new Sequential(9  new Linear(2, 16), new ReLU(), new Linear(16, 1)10);11console.log("Parameters:", Array.from(model.parameters()).length);1213// Training data: y = x0 + 2*x114const X = parameter([[1,0],[0,1],[1,1],[2,1],[1,2],[3,1],[2,2],[0,3]]);15const yTargets = parameter([[1],[2],[3],[4],[5],[5],[6],[6]]);1617const optimizer = new Adam(model.parameters(), { lr: 0.01 });1819// Training loop20for (let epoch = 0; epoch < 200; epoch++) {21  const pred = model.forward(X);22  const diff = pred.sub(yTargets);23  const loss = diff.mul(diff).mean();2425  optimizer.zeroGrad();26  loss.backward();27  optimizer.step();2829  if (epoch % 50 === 0) {30    console.log(`Epoch ${epoch}: loss = ${Number(loss.tensor.data[0]).toFixed(6)}`);31  }32}3334// Inference (plain tensors, no gradient tracking)35const finalPred = model.forward(X.tensor);36console.log("\nPredictions:", finalPred.toString());37console.log("Targets:    ", yTargets.tensor.toString());

Console Output

$ npx tsx 13-neural-network-training/index.ts
=== Neural Network Training ===

Parameters: 4
Epoch 0:   loss = 12.345678
Epoch 50:  loss = 0.234567
Epoch 100: loss = 0.045678
Epoch 150: loss = 0.012345

Predictions: [[1.02],[1.98],[3.01],[3.99],[5.02],[5.01],[5.98],[6.01]]
Targets:     [[1],[2],[3],[4],[5],[5],[6],[6]]
PreviousNext