All Examples
Example 32
advanced
32
Module
Serialization
Neural Networks

Neural Network Module System

The Module base class is the foundation of Deepbox's neural network system, the foundation of Deepbox's neural network system. This example demonstrates building custom Module subclasses with explicit parameter registration via registerModule(), enumerating all parameters with .parameters(), serializing model weights with .stateDict() and restoring with .loadStateDict(), switching between .train() and .eval() modes (which affect Dropout and BatchNorm behavior), freezing parameters with .freezeParameters() to disable gradient tracking (useful for transfer learning), and using the Sequential container for simple linear stacks. You will see the full lifecycle of a neural network module from construction to serialization.

Deepbox Modules Used

deepbox/ndarraydeepbox/nn

What You Will Learn

  • Extend Module and call registerModule() for sub-layers
  • .parameters() yields all learnable tensors recursively
  • stateDict()/loadStateDict() for model serialization and checkpointing
  • .train()/.eval() toggle mode — affects Dropout and BatchNorm
  • .freezeParameters() stops gradient tracking — for transfer learning

Source Code

32-module-system/index.ts
1import { GradTensor, parameter, tensor } from "deepbox/ndarray";2import { Linear, Module, ReLU, Sequential } from "deepbox/nn";34console.log("=== Module System ===\n");56// Custom Module7class MyNet extends Module {8  fc1: Linear; relu: ReLU; fc2: Linear;910  constructor(input: number, hidden: number, output: number) {11    super();12    this.fc1 = new Linear(input, hidden);13    this.relu = new ReLU();14    this.fc2 = new Linear(hidden, output);15    this.registerModule("fc1", this.fc1);16    this.registerModule("relu", this.relu);17    this.registerModule("fc2", this.fc2);18  }1920  override forward(x: Tensor | GradTensor) {21    let out = this.fc1.forward(x);22    out = this.relu.forward(out);23    return this.fc2.forward(out);24  }25}2627const net = new MyNet(4, 8, 2);28console.log("Parameters:", Array.from(net.parameters()).length);2930// State dict31const state = net.stateDict();32console.log("State keys:", Object.keys(state.parameters));33net.loadStateDict(state);34console.log("State loaded ✓");3536// Train/eval mode37net.train();38console.log("Training:", net.training);39net.eval();40console.log("Training:", net.training);4142// Freeze/unfreeze43net.freezeParameters();44console.log("Frozen: all grads disabled");45net.unfreezeParameters();46console.log("Unfrozen: grads re-enabled");

Console Output

$ npx tsx 32-module-system/index.ts
=== Module System ===

Parameters: 4
State keys: ["fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias"]
State loaded ✓
Training: true
Training: false
Frozen: all grads disabled
Unfrozen: grads re-enabled
PreviousProjects