Neural Network Training

1/**2 * Example 13: Neural Network Training3 *4 * Build and train neural networks using the nn, optim, and autograd modules.5 * Covers Sequential models, custom modules, loss functions, and optimizers.6 */78import { isNumericTypedArray, isTypedArray } from "deepbox/core";9import { GradTensor, parameter, type Tensor, tensor } from "deepbox/ndarray";10import { Linear, Module, mseLoss, ReLU, Sequential } from "deepbox/nn";11import { Adam, SGD } from "deepbox/optim";1213console.log("=== Neural Network Training ===\n");1415// Helper to read a scalar value from tensor data16const scalarValue = (t: Tensor): number => {17  const d = t.data;18  if (!isTypedArray(d) || !isNumericTypedArray(d)) return NaN;19  return Number(d[t.offset] ?? 0);20};2122// ---------------------------------------------------------------------------23// Part 1: Sequential model with autograd training24// ---------------------------------------------------------------------------25console.log("--- Part 1: Sequential Model with Autograd ---");2627const model = new Sequential(new Linear(2, 16), new ReLU(), new Linear(16, 1));2829const paramCount = Array.from(model.parameters()).length;30console.log("Model parameters:", paramCount);3132// Training data: y = x0 + 2*x133const X = parameter([34  [1, 0],35  [0, 1],36  [1, 1],37  [2, 1],38  [1, 2],39  [3, 1],40  [2, 2],41  [0, 3],42]);43const yTargets = parameter([[1], [2], [3], [4], [5], [5], [6], [6]]);4445const optimizer = new Adam(model.parameters(), { lr: 0.01 });4647console.log("Training for 200 epochs...");48for (let epoch = 0; epoch < 200; epoch++) {49  // Forward pass with GradTensor builds the computation graph50  const pred = model.forward(X);5152  // Compute MSE loss using GradTensor ops (tracks gradients)53  if (!(pred instanceof GradTensor)) throw new Error("Expected GradTensor from forward");54  const diff = pred.sub(yTargets);55  const loss = diff.mul(diff).mean();5657  // Backward pass and optimize58  optimizer.zeroGrad();59  loss.backward();60  optimizer.step();6162  if (epoch % 50 === 0) {63    console.log(`  Epoch ${epoch}: loss = ${scalarValue(loss.tensor).toFixed(6)}`);64  }65}6667// Evaluate using plain Tensor forward pass (no gradient tracking)68const finalPred = model.forward(X.tensor);69console.log("Predictions:", finalPred.toString());70console.log("Targets:    ", yTargets.tensor.toString());7172// ---------------------------------------------------------------------------73// Part 2: Custom Module74// ---------------------------------------------------------------------------75console.log("\n--- Part 2: Custom Module ---");7677class TwoLayerNet extends Module {78  fc1: Linear;79  relu: ReLU;80  fc2: Linear;8182  constructor(inputDim: number, hiddenDim: number, outputDim: number) {83    super();84    this.fc1 = new Linear(inputDim, hiddenDim);85    this.relu = new ReLU();86    this.fc2 = new Linear(hiddenDim, outputDim);87    this.registerModule("fc1", this.fc1);88    this.registerModule("relu", this.relu);89    this.registerModule("fc2", this.fc2);90  }9192  override forward(x: GradTensor): GradTensor;93  override forward(x: Tensor): Tensor;94  override forward(x: Tensor | GradTensor): Tensor | GradTensor {95    if (x instanceof GradTensor) {96      let out: GradTensor = this.fc1.forward(x);97      out = this.relu.forward(out);98      return this.fc2.forward(out);99    }100    let out: Tensor = this.fc1.forward(x);101    out = this.relu.forward(out);102    return this.fc2.forward(out);103  }104}105106const net = new TwoLayerNet(2, 8, 1);107const netParamCount = Array.from(net.parameters()).length;108console.log("Custom module parameters:", netParamCount);109110// Train/eval mode111net.train();112console.log("Training mode:", net.training);113net.eval();114console.log("Eval mode:", net.training);115116// State dict for serialization117const state = net.stateDict();118console.log("State dict keys:", Object.keys(state.parameters).join(", "));119120// ---------------------------------------------------------------------------121// Part 3: Plain Tensor forward pass with mseLoss122// ---------------------------------------------------------------------------123console.log("\n--- Part 3: Plain Tensor Forward + mseLoss ---");124125const inputTensor = tensor([126  [1, 0],127  [0, 1],128  [1, 1],129  [2, 1],130]);131const targetTensor = tensor([[1], [2], [3], [4]]);132133// Forward pass with plain Tensors — no autograd, just inference134const rawOutput = model.forward(inputTensor);135const output = rawOutput instanceof GradTensor ? rawOutput.tensor : rawOutput;136const evalLoss = mseLoss(output, targetTensor);137console.log("Eval loss (plain Tensor):", scalarValue(evalLoss).toFixed(6));138139// ---------------------------------------------------------------------------140// Part 4: SGD with momentum141// ---------------------------------------------------------------------------142console.log("\n--- Part 4: SGD with Momentum ---");143144const sgdModel = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));145const sgdOptimizer = new SGD(sgdModel.parameters(), {146  lr: 0.01,147  momentum: 0.9,148});149150for (let epoch = 0; epoch < 100; epoch++) {151  const pred = sgdModel.forward(X);152  if (!(pred instanceof GradTensor)) throw new Error("Expected GradTensor from forward");153  const diff = pred.sub(yTargets);154  const loss = diff.mul(diff).mean();155  sgdOptimizer.zeroGrad();156  loss.backward();157  sgdOptimizer.step();158}159160const rawSgdPred = sgdModel.forward(X.tensor);161const sgdPred = rawSgdPred instanceof GradTensor ? rawSgdPred.tensor : rawSgdPred;162const sgdLoss = mseLoss(sgdPred, yTargets.tensor);163console.log("SGD final loss:", scalarValue(sgdLoss).toFixed(6));164165console.log("\n=== Neural Network Training Complete ===");166