Neural Network Module System

1/**2 * Example 32: Neural Network Module System3 *4 * The Module base class: parameter registration, state dicts, train/eval5 * modes, freeze/unfreeze, forward hooks, and the Sequential container.6 */7 8import { type AnyTensor, GradTensor, noGrad, tensor } from "deepbox/ndarray";9import { Linear, Module, ReLU, Sequential } from "deepbox/nn";10 11console.log("=== Neural Network Module System ===\n");12 13// ---------------------------------------------------------------------------14// Part 1: Custom Module with parameter registration15// ---------------------------------------------------------------------------16console.log("--- Part 1: Custom Module ---");17 18class MyNet extends Module {19  fc1: Linear;20  relu: ReLU;21  fc2: Linear;22 23  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  }32 33  override forward(x: AnyTensor): AnyTensor {34    return this.fc2.forward(this.relu.forward(this.fc1.forward(x)));35  }36}37 38const net = new MyNet(4, 8, 2);39console.log("MyNet(4 -> 8 -> 2)");40 41// ---------------------------------------------------------------------------42// Part 2: Parameter enumeration43// ---------------------------------------------------------------------------44console.log("\n--- Part 2: Parameters ---");45 46// parameters() and namedParameters() yield GradTensors, the type optimizers take47const named = Array.from(net.namedParameters());48console.log(`Total parameter tensors: ${named.length}`);49for (const [name, p] of named) {50  console.log(`  ${name.padEnd(11)} shape [${p.shape.join(", ")}]`);51}52 53// ---------------------------------------------------------------------------54// Part 3: State dict (serialization & loading)55// ---------------------------------------------------------------------------56console.log("\n--- Part 3: State Dict ---");57 58const stateDict = net.stateDict();59console.log("State dict parameter keys:");60for (const key of Object.keys(stateDict.parameters)) {61  console.log(`  ${key}`);62}63 64// Load the state dict back, for example from a saved checkpoint65net.loadStateDict(stateDict);66console.log("State dict loaded");67 68// ---------------------------------------------------------------------------69// Part 4: Train/Eval mode70// ---------------------------------------------------------------------------71console.log("\n--- Part 4: Train/Eval Mode ---");72 73net.train();74console.log(`After train(): training = ${net.training}`);75 76net.eval();77console.log(`After eval():  training = ${net.training}`);78console.log("  eval() turns dropout off and makes batch norm use its running statistics");79console.log("  It does not stop gradient tracking. Use noGrad() for that.");80 81// ---------------------------------------------------------------------------82// Part 5: Freeze/Unfreeze parameters83// ---------------------------------------------------------------------------84console.log("\n--- Part 5: Freeze/Unfreeze ---");85 86const countTrainable = (module: Module): number =>87  Array.from(module.parameters()).filter((p) => p.requiresGrad).length;88 89console.log(`Trainable parameter tensors: ${countTrainable(net)}`);90 91net.freezeParameters();92console.log(`After freezeParameters(): ${countTrainable(net)}`);93 94// Freeze or unfreeze selected parameters by name95net.unfreezeParameters(["fc2.weight", "fc2.bias"]);96console.log(`After unfreezeParameters(["fc2.weight", "fc2.bias"]): ${countTrainable(net)}`);97 98net.unfreezeParameters();99console.log(`After unfreezeParameters(): ${countTrainable(net)}`);100 101// ---------------------------------------------------------------------------102// Part 6: Sequential container103// ---------------------------------------------------------------------------104console.log("\n--- Part 6: Sequential Container ---");105 106const seqModel = new Sequential(new Linear(4, 8), new ReLU(), new Linear(8, 2));107 108console.log("Sequential(Linear(4,8), ReLU, Linear(8,2))");109console.log(`Parameter tensors: ${Array.from(seqModel.parameters()).length}`);110 111// A plain tensor goes in, no parameter(...) wrapping needed112const input = tensor([[1, 2, 3, 4]]);113 114// Training: the output is a GradTensor that tracks the weights115const trainOutput = seqModel.forward(input);116console.log(`Input shape:  [${input.shape.join(", ")}]`);117console.log(`Output shape: [${trainOutput.shape.join(", ")}]`);118console.log(`Output is a GradTensor: ${GradTensor.isGradTensor(trainOutput)}`);119console.log(120  `Output requiresGrad: ${GradTensor.isGradTensor(trainOutput) && trainOutput.requiresGrad}`121);122 123// Inference: inside noGrad() no graph is built and the output is a plain tensor124const inferOutput = noGrad(() => seqModel.forward(input));125console.log(`Inside noGrad(), output is a GradTensor: ${GradTensor.isGradTensor(inferOutput)}`);126 127// A model with every parameter frozen also returns a plain tensor128seqModel.freezeParameters();129console.log(130  `With frozen parameters, output is a GradTensor: ${GradTensor.isGradTensor(seqModel.forward(input))}`131);132seqModel.unfreezeParameters();133 134// ---------------------------------------------------------------------------135// Part 7: Forward hooks136// ---------------------------------------------------------------------------137console.log("\n--- Part 7: Forward Hooks ---");138 139// call() runs the hooks around forward(). Calling forward() directly skips them.140const removeHook = seqModel.registerForwardHook((_module, _inputs, output) => {141  console.log(`  hook saw output shape [${output.shape.join(", ")}]`);142  return undefined; // keep the output unchanged143});144noGrad(() => seqModel.call(input));145removeHook();146console.log("Hook removed: the next call prints nothing");147noGrad(() => seqModel.call(input));148 149console.log("\n=== Module System Complete ===");