Neural Network Module System

1/**2 * Example 32: Neural Network Module System3 *4 * Demonstrates the Module base class features: parameter registration, state5 * serialization, train/eval modes, freeze/unfreeze, and forward hooks.6 */78import { GradTensor, parameter, type Tensor, tensor } from "deepbox/ndarray";9import { Linear, Module, ReLU, Sequential } from "deepbox/nn";1011console.log("=== Neural Network Module System ===\n");1213// ---------------------------------------------------------------------------14// Part 1: Custom Module with parameter registration15// ---------------------------------------------------------------------------16console.log("--- Part 1: Custom Module ---");1718class MyNet extends Module {19  fc1: Linear;20  relu: ReLU;21  fc2: Linear;2223  constructor(inputDim: number, hiddenDim: number, outputDim: number) {24    super();25    this.fc1 = new Linear(inputDim, hiddenDim);26    this.relu = new ReLU();27    this.fc2 = new Linear(hiddenDim, outputDim);28    this.registerModule("fc1", this.fc1);29    this.registerModule("relu", this.relu);30    this.registerModule("fc2", this.fc2);31  }3233  override forward(x: GradTensor): GradTensor;34  override forward(x: Tensor): Tensor;35  override forward(x: Tensor | GradTensor): Tensor | GradTensor {36    if (x instanceof GradTensor) {37      let out: GradTensor = this.fc1.forward(x);38      out = this.relu.forward(out);39      return this.fc2.forward(out);40    }41    let out: Tensor = this.fc1.forward(x);42    out = this.relu.forward(out);43    return this.fc2.forward(out);44  }45}4647const net = new MyNet(4, 8, 2);48console.log("MyNet(4 -> 8 -> 2)");4950// ---------------------------------------------------------------------------51// Part 2: Parameter enumeration52// ---------------------------------------------------------------------------53console.log("\n--- Part 2: Parameters ---");5455const params = Array.from(net.parameters());56console.log(`Total parameter tensors: ${params.length}`);57for (const p of params) {58  const t = p instanceof GradTensor ? p.tensor : p;59  console.log(`  Shape: [${t.shape.join(", ")}]`);60}6162// ---------------------------------------------------------------------------63// Part 3: State dict — serialization & loading64// ---------------------------------------------------------------------------65console.log("\n--- Part 3: State Dict ---");6667const stateDict = net.stateDict();68console.log("State dict parameter keys:");69for (const key of Object.keys(stateDict.parameters)) {70  console.log(`  ${key}`);71}7273// Load state dict back (e.g., from a saved checkpoint)74net.loadStateDict(stateDict);75console.log("State dict loaded successfully");7677// ---------------------------------------------------------------------------78// Part 4: Train/Eval mode79// ---------------------------------------------------------------------------80console.log("\n--- Part 4: Train/Eval Mode ---");8182net.train();83console.log(`Training mode: ${net.training}`);8485net.eval();86console.log(`Eval mode: ${net.training}`);87console.log("  Eval mode disables dropout and uses running stats for batchnorm");8889// ---------------------------------------------------------------------------90// Part 5: Freeze/Unfreeze parameters91// ---------------------------------------------------------------------------92console.log("\n--- Part 5: Freeze/Unfreeze ---");9394net.freezeParameters();95console.log("After freezeParameters:");96const frozenParams = Array.from(net.parameters());97const frozenGrads = frozenParams.filter((p) => p instanceof GradTensor && p.requiresGrad);98console.log(`  Parameters requiring grad: ${frozenGrads.length}`);99100net.unfreezeParameters();101console.log("After unfreezeParameters:");102const unfrozenParams = Array.from(net.parameters());103const unfrozenGrads = unfrozenParams.filter((p) => p instanceof GradTensor && p.requiresGrad);104console.log(`  Parameters requiring grad: ${unfrozenGrads.length}`);105106// ---------------------------------------------------------------------------107// Part 6: Sequential container108// ---------------------------------------------------------------------------109console.log("\n--- Part 6: Sequential Container ---");110111const seqModel = new Sequential(new Linear(4, 8), new ReLU(), new Linear(8, 2));112113console.log("Sequential(Linear(4,8), ReLU, Linear(8,2))");114const seqParams = Array.from(seqModel.parameters()).length;115console.log(`Parameters: ${seqParams}`);116117// Forward pass with plain Tensor (inference)118const input = tensor([[1, 2, 3, 4]]);119const output = seqModel.forward(input);120const outTensor = output instanceof GradTensor ? output.tensor : output;121console.log(`Input shape:  [${input.shape.join(", ")}]`);122console.log(`Output shape: [${outTensor.shape.join(", ")}]`);123124// Forward pass with GradTensor (training)125const gradInput = parameter([[1, 2, 3, 4]]);126const gradOutput = seqModel.forward(gradInput);127console.log(128  `GradTensor output requiresGrad: ${gradOutput instanceof GradTensor ? gradOutput.requiresGrad : false}`129);130131console.log("\n=== Module System Complete ===");132