deepbox/nn

Module, Containers & Trainer

The Module base class, containers, gradient clipping, Trainer, and the early-stopping and checkpoint helpers.

Training

Overview

model.forward(x) on a plain tensor returns a GradTensor that tracks the weights when gradient mode is on and the module has trainable parameters, so training data does not need parameter(). Inside noGrad() it returns a plain Tensor. model.eval() only changes layers such as dropout and batch normalization, and it does not stop tracking, so wrap inference in noGrad. Modules start in training mode.
A training step is optimizer.zeroGrad(), forward, a loss, loss.backward(), optimizer.step(). Trainer runs that loop over batches of [input, target] pairs from an array or a DataLoader. Its options are epochs (default 10), earlyStopping, maxGradNorm, accumulationSteps and restoreBestWeights, which reloads the weights of the best epoch. The loss function must return a single-element loss. fitAsync accepts asynchronous batch sources.
await model.to("webgpu") moves float32, float16 and bfloat16 parameters and buffers to the GPU once a WebGPU backend is registered. to returns a promise, and without a registered backend it throws DeviceError. Embedding, EmbeddingBag, Conv3d, ConvTranspose1d, ConvTranspose2d, the recurrent layers, PReLU and SpectralNorm keep their weights on the host, and those layers keep working after to.
On a device, autograd and the optimizer step run on the device for the operations listed in Devices & Execution. Trainer and clipGradNorm_ read values on the host, so they are meant for CPU tensors. For GPU training, write the loop yourself and read the loss with await loss.cpu().
freezeParameters(names?) turns off requiresGrad, and optimizers skip frozen parameters. stateDict() and loadStateDict() save and restore parameters and buffers. EarlyStopping defaults to patience: 10, while Trainer's earlyStopping option defaults to patience: 5.

Types

ForwardHook

type

Hook function called after the forward pass.

type ForwardHook = (module: Module, inputs: AnyTensor[], output: AnyTensor) => AnyTensor | undefined;

ForwardPreHook

type

Hook function called before the forward pass.

type ForwardPreHook = (module: Module, inputs: AnyTensor[]) => AnyTensor[] | undefined;

EpochInfo

type

Information about a completed epoch.

readonly epoch: number
readonly trainLoss: number
Mean training loss over the epoch, weighted by the first dimension of each input batch.
readonly valLoss: number | undefined
Mean validation loss, or undefined when no validation data was given.

LossFn

type

A loss function that takes model output and targets, returning a scalar loss. The output may be a Tensor or GradTensor (AnyTensor) depending on training mode.

type LossFn = (output: AnyTensor, target: Tensor) => AnyTensor;

TrainerCallback

type

A callback invoked at the end of each epoch.

type TrainerCallback = (info: EpochInfo) => void;

TrainerOptimizer

interface

An optimizer-like object with step() and zeroGrad() methods.

step(): void
zeroGrad(): void

TrainerOptions

type

Options for the Trainer.

readonly epochs?: number
Number of training epochs.
readonly earlyStopping?: { readonly patience?: number; readonly minDelta?: number; }
Early stopping configuration (monitors the validation loss, or the training loss when no validation data is given).
readonly checkpoint?: ModelCheckpoint
Model checkpointing helper.
readonly callbacks?: readonly TrainerCallback[]
Callbacks invoked after each epoch, including the one that triggers early stopping.
readonly verbose?: boolean
Whether to log epoch progress to the console.
readonly maxGradNorm?: number
Clip the global L2 norm of the gradients to this value before every optimizer step.
readonly accumulationSteps?: number
Accumulate the gradients of this many batches before every optimizer step (default: 1).
readonly restoreBestWeights?: boolean
After training, load the weights of the epoch with the lowest monitored loss (the validation loss, or the training loss when there is no validation data) back into the model (default: false).

TrainerResult

type

Result returned by Trainer.fit.

readonly history: EpochInfo[]
readonly stoppedEarly: boolean
readonly bestEpoch: number | undefined
1-based epoch with the best monitored loss.

Classes

ModuleDict

extends Module

Holds submodules in a dictionary.

ModuleDict can be indexed by string keys and the modules it contains are properly registered.

ExampleTypeScript

const modules = new ModuleDict({ encoder: new Linear(10, 5), decoder: new Linear(5, 10) });const enc = modules.get('encoder');

Members (15)

constructor(modules?: Record<string, Module>)
length: number
Number of modules in the dict.
get(key: string): Module
Get a module by key.
has(key: string): boolean
Check if key exists.
set(key: string, module: Module): this
Set a module by key, replacing an existing entry in place.
update(modules: Record<string, Module> | Iterable<readonly [ string, Module ]>): this
Add every [key, module] pair from a record, a Map or an iterable of pairs.
delete(key: string): boolean
Delete a module by key.
pop(key: string): Module
Remove a module by key and return it.
clear(): void
Remove all modules.
keys(): IterableIterator<string>
Get all keys.
values(): IterableIterator<Module>
Get all values.
entries(): IterableIterator<[ string, Module ]>
Get all entries.
[Symbol.iterator](): Iterator<[ string, Module ]>
Iterate over entries.
forward(..._inputs: AnyTensor[]): AnyTensor
toString(): string

ModuleList

extends Module

Holds submodules in a list.

ModuleList can be indexed like a regular array and the modules it contains are properly registered so that all Module methods (parameters, to, etc.) work.

ExampleTypeScript

const layers = new ModuleList([new Linear(10, 5), new ReLU()]);layers.append(new Linear(5, 2));for (const layer of layers) { ... }

Members (10)

constructor(modules?: Iterable<Module>)
length: number
Number of modules in the list.
get(index: number): Module
Get a module by index.
append(module: Module): this
Append a module to the end.
extend(modules: Iterable<Module>): this
Append every module of an iterable, in order.
insert(index: number, module: Module): this
Insert a module at a given index; later modules shift one position up.
pop(index?: number): Module
Remove and return the module at index (default: the last one).
[Symbol.iterator](): Iterator<Module>
Iterate over modules.
forward(..._inputs: AnyTensor[]): AnyTensor
toString(): string

ParameterDict

extends Module

Holds parameters in a dictionary.

ParameterDict can be indexed by string keys and the parameters it contains are properly registered so that all Module methods work correctly.

ExampleTypeScript

const params = new ParameterDict({ weight: weightParam, bias: biasParam });const w = params.get('weight');

Members (15)

constructor(parameters?: Record<string, GradTensor>)
length: number
Number of parameters in the dict.
get(key: string): GradTensor
Get a parameter by key.
has(key: string): boolean
Check if key exists.
set(key: string, param: GradTensor): this
Set a parameter by key, replacing an existing entry in place.
update(parameters: Record<string, GradTensor> | Iterable<readonly [ string, GradTensor ]>): this
Add every [key, parameter] pair from a record, a Map or an iterable of pairs.
delete(key: string): boolean
Delete a parameter by key.
pop(key: string): GradTensor
Remove a parameter by key and return it.
clear(): void
Remove all parameters.
keys(): IterableIterator<string>
Get all keys.
values(): IterableIterator<GradTensor>
Get all values.
entries(): IterableIterator<[ string, GradTensor ]>
Get all entries.
[Symbol.iterator](): Iterator<[ string, GradTensor ]>
Iterate over entries.
forward(..._inputs: AnyTensor[]): AnyTensor
toString(): string

ParameterList

extends Module

Holds parameters in a list.

ParameterList can be indexed like a regular array and the parameters it contains are properly registered so that all Module methods (parameters, namedParameters, etc.) work correctly.

ExampleTypeScript

const params = new ParameterList([param1, param2]);params.append(param3);for (const p of params) { ... }

Members (8)

constructor(parameters?: Iterable<GradTensor>)
length: number
Number of parameters in the list.
get(index: number): GradTensor
Get a parameter by index.
append(param: GradTensor): this
Append a parameter to the end.
extend(parameters: Iterable<GradTensor>): this
Append every parameter of an iterable, in order.
[Symbol.iterator](): Iterator<GradTensor>
Iterate over parameters.
forward(..._inputs: AnyTensor[]): AnyTensor
toString(): string

Sequential

extends Module

Sequential container for stacking layers in a linear pipeline.

Purpose: - Simplifies model construction by chaining layers sequentially - Automatically manages forward pass through all layers - Provides clean API for building feedforward networks

Behavior: The output of each layer becomes the input to the next layer. Layers are executed in the order they were added.

A plain Tensor input needs no wrapping for training: the first layer with trainable weights returns a GradTensor that tracks them, and the layers after it pass it on. Inside noGrad(), or when every weight is frozen, the result is a plain Tensor.

ExampleTypeScript

import { Sequential, Linear, ReLU, Dropout } from 'deepbox/nn';import { tensor } from 'deepbox/ndarray'; // Create a simple feedforward networkconst model = new Sequential(  new Linear(784, 256),  new ReLU(),  new Dropout(0.5),  new Linear(256, 10)); const input = tensor(new Array(784).fill(0));const output = model.forward(input);

ExampleTypeScript

// Access individual layersconst model = new Sequential(  new Linear(10, 5),  new ReLU()); const firstLayer = model.getLayer(0); // Linear layerconst layerCount = model.length; // 2

Members (9)

constructor(...layers: Module[])
Create a new Sequential container.
length: number
Get the number of layers in the sequential container.
forward(...inputs: AnyTensor[]): AnyTensor
Forward pass: sequentially apply all layers.
getLayer(index: number): Module
Get a layer by index.
append(layer: Module): this
Append a layer to the end of the pipeline.
extend(layers: Iterable<Module>): this
Append several layers, in order.
insert(index: number, layer: Module): this
Insert a layer at index; later layers shift one position up and their parameter names are renumbered.
toString(): string
Get string representation showing all layers.
[Symbol.iterator](): Iterator<Module>
Iterate over all layers.

Module

Base class for all neural network modules.

All models should subclass this class. Modules can contain other modules, allowing to nest them in a tree structure.

ExampleTypeScript

import { Module, Linear, ReLU } from 'deepbox/nn';import type { AnyTensor, Tensor } from 'deepbox/ndarray'; class MyModel extends Module {  private fc1: Linear;  private relu: ReLU;  private fc2: Linear;   constructor() {    super();    this.fc1 = new Linear(10, 5);    this.relu = new ReLU();    this.fc2 = new Linear(5, 2);    this.registerModule('fc1', this.fc1);    this.registerModule('relu', this.relu);    this.registerModule('fc2', this.fc2);  }   forward(x: Tensor): AnyTensor {    let out = this.fc1.forward(x);    out = this.relu.forward(out);    out = this.fc2.forward(out);    return out;  }}

Members (24)

training: boolean
Check if the module is in training mode.
abstract forward(...inputs: AnyTensor[]): AnyTensor
Forward pass of the module.
call(...inputs: AnyTensor[]): AnyTensor
Run the module: forward pre-hooks, then Module.forward, then forward hooks.
parameters(recurse?: boolean): Generator<GradTensor>
Get all parameters of this module and its children.
namedParameters(prefix?: string, recurse?: boolean, removeDuplicate?: boolean): Generator<[ string, GradTensor ]>
Get all named parameters of this module and its children.
modules(recurse?: boolean): Generator<Module>
Get all child modules.
namedModules(prefix?: string, recurse?: boolean, removeDuplicate?: boolean): Generator<[ string, Module ]>
Get all named child modules.
children(): Generator<Module>
Get the immediate child modules (not the module itself, no recursion).
namedChildren(): Generator<[ string, Module ]>
Get the immediate child modules together with their registered names.
train(mode?: boolean): this
Set the module in training mode.
eval(): this
Set the module in evaluation mode.
zeroGrad(): void
Zero out the gradients of all parameters.
buffers(recurse?: boolean): Generator<Tensor>
Get all buffers of this module and its children.
namedBuffers(prefix?: string, recurse?: boolean, removeDuplicate?: boolean): Generator<[ string, Tensor ]>
Get all named buffers of this module and its children.
freezeParameters(names?: string[], recurse?: boolean): void
Freeze specific parameters by name (or all if none provided).
unfreezeParameters(names?: string[], recurse?: boolean): void
Unfreeze specific parameters by name (or all if none provided).
stateDict(): { parameters: Record<string, StateEntry>; buffers: Record<string, StateEntry>; }
Get the state dictionary of the module.
loadStateDict(stateDict: { parameters?: Record<string, StateEntry>; buffers?: Record<string, StateEntry>; }): void
Load state dictionary into the module.
to(device: Device): Promise<this>
Move the module's parameters and buffers to a device.
apply(fn: (module: Module) => void): this
Apply a function to this module and every descendant, parents before children.
registerForwardPreHook(hook: ForwardPreHook): () => void
Register a forward pre-hook, run by Module.call before forward.
registerForwardHook(hook: ForwardHook): () => void
Register a forward hook, run by Module.call after forward.
toString(): string
Get string representation of the module.
summary(): string
Build a text summary of the model: one row per own parameter and per child module with its parameter count, followed by total, trainable and non-trainable counts.

Trainer

High-level training loop for neural network modules.

Each epoch trains on every training batch (zero gradients, forward, loss, backward, optional gradient clipping, optimizer step; with accumulationSteps the gradients of several batches are summed before each step), then evaluates the validation batches without recording an autograd graph. The loss function must return a single-element loss. Input batches are wrapped in a GradTensor during training so that the model's parameters receive gradients. If the loss is a plain Tensor rather than a GradTensor, no backward pass is run.

fit() can be called again to continue training: the history restarts at epoch 1 and early-stopping state is reset. The model is left in the mode it was last in (evaluation mode when validation data was given).

The loss value is read on the host after each step, so Trainer is meant for CPU tensors. A loss that lives in device memory throws a DeviceError; train on a device with an explicit loop instead.

ExampleTypeScript

import { Linear, mseLoss, Trainer } from 'deepbox/nn';import { Adam } from 'deepbox/optim'; const model = new Linear(2, 1);const optimizer = new Adam(model.parameters(), { lr: 0.001 });const lossFn = (pred, target) => mseLoss(pred, target); const trainer = new Trainer(model, optimizer, lossFn, {  epochs: 50,  earlyStopping: { patience: 5 },  verbose: true,}); const trainData = [[xBatch1, yBatch1], [xBatch2, yBatch2]];const result = trainer.fit(trainData); // Step once per 4 batches and load the best epoch's weights at the endconst accumulating = new Trainer(model, optimizer, lossFn, {  epochs: 50,  accumulationSteps: 4,  restoreBestWeights: true,});

Members (3)

constructor(model: Module, optimizer: TrainerOptimizer, lossFn: LossFn, options?: TrainerOptions)
fit(trainData: Iterable<readonly [ Tensor, Tensor ]>, valData?: Iterable<readonly [ Tensor, Tensor ]>): TrainerResult
Run the training loop.
fitAsync(trainData: AsyncIterable<readonly [ Tensor, Tensor ]> | Iterable<readonly [ Tensor, Tensor ]>, valData?: AsyncIterable<readonly [ Tensor, Tensor ]> | Iterable<readonly [ Tensor, Tensor ]>): Promise<TrainerResult>
Run the training loop over an asynchronous (or synchronous) batch source.

EarlyStopping

Early stopping callback to terminate training when a monitored metric stops improving.

Tracks a metric value across epochs and signals when training should stop based on patience (number of epochs with no improvement).

ExampleTypeScript

import { EarlyStopping } from 'deepbox/nn'; const es = new EarlyStopping({ patience: 5, minDelta: 0.001 }); for (let epoch = 0; epoch < maxEpochs; epoch++) {  const valLoss = trainOneEpoch();  if (es.step(valLoss)) {    console.log(`Early stopping at epoch ${epoch}`);    break;  }}

Members (7)

constructor(options?: { readonly patience?: number; readonly minDelta?: number; readonly mode?: "min" | "max"; })
Create an EarlyStopping callback.
isStopped: boolean
Whether early stopping has been triggered.
best: number
The best metric value observed so far.
bestEpochNum: number
The epoch at which the best metric was observed (1-indexed).
waitCount: number
Number of epochs since last improvement.
step(metric: number): boolean
Report the current metric value and check if training should stop.
reset(): void
Reset the callback state to start fresh.

GradientAccumulator

Gradient accumulation helper for training with effective batch sizes larger than what fits in memory.

Accumulates gradients over multiple forward/backward passes before performing a single optimizer step, then zeros the gradients.

ExampleTypeScript

import { GradientAccumulator } from 'deepbox/nn'; const accumulator = new GradientAccumulator(4); // accumulate over 4 steps for (const batch of dataLoader) {  const loss = forward(batch);  backward(loss);  console.log(`loss ${loss.item()}`);   if (accumulator.step()) {    // Time to update: 4 mini-batches accumulated    optimizer.step();    optimizer.zeroGrad();  }}

Members (8)

constructor(accumSteps: number)
Create a GradientAccumulator.
accumulated: number
Number of mini-batches accumulated since the last optimizer step.
totalProcessed: number
Total number of mini-batches processed.
optimizerSteps: number
Number of optimizer steps that have been triggered.
steps: number
The configured number of accumulation steps.
scaleFactor: number
The scaling factor for loss to account for accumulation.
step(): boolean
Record one mini-batch and check if it's time to perform an optimizer step.
reset(): void
Reset the accumulator state.

ModelCheckpoint

Model checkpoint helper to track and save the best model state during training.

Stores a copy of the model's state dictionary whenever the monitored metric improves.

ExampleTypeScript

import { ModelCheckpoint } from 'deepbox/nn'; const checkpoint = new ModelCheckpoint({ mode: 'min' }); for (let epoch = 0; epoch < maxEpochs; epoch++) {  const valLoss = trainOneEpoch();  checkpoint.step(model, valLoss);} // Restore best modelcheckpoint.restore(model);

Members (7)

constructor(options?: { readonly mode?: "min" | "max"; })
Create a ModelCheckpoint.
hasSavedState: boolean
Whether a checkpoint has been saved.
best: number
The best metric value observed.
bestEpochNum: number
The epoch at which the best metric was observed (1-indexed).
step(model: Module, metric: number): boolean
Check the metric and save model state if improved.
restore(model: Module): void
Restore the best saved model state.
reset(): void
Reset the checkpoint state.

Functions

clipGradNorm

function clipGradNorm(parameters: GradTensor | Iterable<GradTensor>, maxNorm: number, normType?: number, errorIfNonfinite?: boolean): number;

Clip the total norm of gradients of an iterable of parameters.

Details

The norm is computed over all gradients together, as if they were concatenated into a single vector. Gradients are modified in-place, and parameters without a gradient are skipped. A parameter that appears more than once is counted and scaled once.

Like PyTorch, gradients are multiplied by maxNorm / (totalNorm + 1e-6) whenever that factor is below 1. If the norm is NaN the gradients are left unchanged; pass errorIfNonfinite to fail loudly instead.

Gradients are read on the host, in logical order, so strided views and a storage offset are handled. Gradients stored on a kernel device (for example webgpu) cannot be read synchronously: the call throws a DeviceError before any gradient is changed. Move them back with await grad.cpu() and setGrad first.

Parameters

parameters GradTensor | Iterable<GradTensor>: A GradTensor or an iterable of GradTensors whose gradients will be clipped
maxNorm number: Maximum allowed norm value, must be >= 0
normTypeoptional number: Type of norm: Infinity for the max norm, otherwise the p > 0 of the p-norm (default: 2 for the L2 norm)
errorIfNonfiniteoptional boolean: Throw if the total norm is NaN or infinite (default: false)

Returns

numberThe total norm of the gradients (before clipping)

Throws

InvalidParameterErrorIf maxNorm is negative or NaN, or normType is not positive
DataValidationErrorIf errorIfNonfinite is set and the total norm is not finite
DTypeErrorIf a gradient is not floating-point
DeviceErrorIf a gradient lives in device memory

clipGradNorm_

function clipGradNorm_(parameters: GradTensor | Iterable<GradTensor>, maxNorm: number, normType?: number, errorIfNonfinite?: boolean): number;

Clip the total norm of gradients of an iterable of parameters.

Details

Parameters

parameters GradTensor | Iterable<GradTensor>: A GradTensor or an iterable of GradTensors whose gradients will be clipped
maxNorm number: Maximum allowed norm value, must be >= 0
normTypeoptional number: Type of norm: Infinity for the max norm, otherwise the p > 0 of the p-norm (default: 2 for the L2 norm)
errorIfNonfiniteoptional boolean: Throw if the total norm is NaN or infinite (default: false)

Returns

numberThe total norm of the gradients (before clipping)

Throws

InvalidParameterErrorIf maxNorm is negative or NaN, or normType is not positive
DataValidationErrorIf errorIfNonfinite is set and the total norm is not finite
DTypeErrorIf a gradient is not floating-point
DeviceErrorIf a gradient lives in device memory

clipGradValue

function clipGradValue(parameters: GradTensor | Iterable<GradTensor>, clipValue: number): void;

Clip the gradients of an iterable of parameters at specified value.

Each gradient element is clamped to [-clipValue, clipValue]. Gradients are modified in-place. NaN entries stay NaN. A parameter that appears more than once is clipped once. Gradients stored on a kernel device throw a DeviceError, as described for clip_grad_norm_.

Parameters

parameters GradTensor | Iterable<GradTensor>: A GradTensor or an iterable of GradTensors whose gradients will be clipped
clipValue number: Maximum absolute value for gradient elements, must be >= 0

Throws

InvalidParameterErrorIf clipValue is negative or NaN
DTypeErrorIf a gradient is not floating-point
DeviceErrorIf a gradient lives in device memory

clipGradValue_

function clipGradValue_(parameters: GradTensor | Iterable<GradTensor>, clipValue: number): void;

Clip the gradients of an iterable of parameters at specified value.

Parameters

parameters GradTensor | Iterable<GradTensor>: A GradTensor or an iterable of GradTensors whose gradients will be clipped
clipValue number: Maximum absolute value for gradient elements, must be >= 0

Throws

InvalidParameterErrorIf clipValue is negative or NaN
DTypeErrorIf a gradient is not floating-point
DeviceErrorIf a gradient lives in device memory

Deprecated names

These older spellings still work in 1.x. New code should use the names above.

clip_grad_norm_Prefer clipGradNorm_.
clip_grad_value_Prefer clipGradValue_.

nn-module.tsTypeScript

import { noGrad, tensor } from "deepbox/ndarray";import { Linear, ReLU, Sequential, Trainer, mseLoss } from "deepbox/nn";import { Adam } from "deepbox/optim";import { setSeed } from "deepbox/random"; setSeed(0); const X = tensor([[0, 0], [0, 1], [1, 0], [1, 1]]);const y = tensor([[0], [2], [1], [3]]); // y = x1 + 2 * x2 const model = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));const optimizer = new Adam(model.parameters(), { lr: 0.05 });const trainer = new Trainer(model, optimizer, (output, target) => mseLoss(output, target), {  epochs: 200,  maxGradNorm: 5,  restoreBestWeights: true,}); const result = trainer.fit([[X, y]]); // an array of [input, target] batchesconsole.log(result.history.length, result.bestEpoch, result.stoppedEarly); model.eval();const prediction = noGrad(() => model.forward(X)); // plain Tensor, no graphconsole.log(prediction.toString());console.log(Object.keys(model.stateDict()));

Found a mistake? Open an issue.