Wiki
Download
Manual
Eggs
API
Tests
Bugs
show
edit
You can edit this page using
wiki syntax
for markup.
Article contents:
[[tags: egg math ai machine-learning]] [[toc:]] == nanograd A lightweight automatic differentiation and neural network framework for CHICKEN Scheme, with BLAS-accelerated operations, batch processing, YASOS-based object abstractions, and an optional lazy-tensor backend built on [[array-morphisms]] that compiles a whole training step into a replayable program. === Description nanograd provides a framework for building and training neural networks with automatic differentiation. It features: * Reverse-mode automatic differentiation * Batch processing support throughout the stack * BLAS-accelerated linear algebra with batched GEMM * YASOS-based polymorphic object system * 32-bit and 64-bit floating-point tensors * Dense, convolutional, batch normalization, layer normalization, pooling and flatten layers * SGD, Adam and RMSprop optimizers, with an optional fast path for optimizer update kernels written in C * Learning rate schedules and gradient clipping * Gradient health diagnostics * Model checkpoints in CBOR (RFC 8949) format * Lazy tensors backed by [[array-morphisms]]: fused operations, buffer reuse, and training steps compiled to an SSA program that is replayed for every batch === Requirements * [[yasos]] * [[blas]] * [[mathh]] * [[vector-lib]] * [[srfi-1]] * [[srfi-4]] * [[srfi-69]] * [[cbor]] * [[array-morphisms]] The unit tests also need [[test]]. === Modules ; nanograd-autograd : Tensors, operations and the automatic differentiation engine. ; nanograd-layer : Layers, containers, activation objects and checkpoints. ; nanograd-optimizer : SGD, Adam and RMSprop. ; nanograd-training : Learning rate schedules and gradient clipping. ; nanograd-diagnostics : Gradient statistics, health checks and monitoring. ; nanograd-array-morphisms : Lazy tensors, AM-backed layers and compiled training steps. ==== nanograd-autograd Core automatic differentiation engine with tensor operations and batch support. ===== Strict and Lazy Tensors nanograd has two kinds of tensor. Both answer the same tensor operations ({{tensor-data}}, {{tensor-shape}}, {{backward!}}, and so on), so layers and optimizers work with either. ; Strict tensors : Created by {{make-tensor32}} and {{make-tensor64}}. Their data is an SRFI-4 vector, and every operation computes its result at once. ; Lazy tensors : Created by the {{nanograd-array-morphisms}} module. They wrap an array-morphisms variable, and their values are computed only when needed. See the {{nanograd-array-morphisms}} module below. <procedure>(strict-tensor? obj) -> boolean</procedure> <procedure>(lazy-tensor? obj) -> boolean</procedure> Return #t if {{obj}} is a strict or a lazy tensor, respectively. <procedure>(tensor-force tensor) -> tensor</procedure> Returns a strict tensor holding the values of {{tensor}}. A strict tensor is returned unchanged. <procedure>(tensor-morphism-value tensor) -> morphism or #f</procedure> Returns the array morphism behind a lazy tensor, or #f for a strict tensor. <procedure>(lazy-tensor-mv tensor) -> morph-variable</procedure> Returns the array-morphisms variable of a lazy tensor. ===== Tensor Constructors <procedure>(make-tensor32 data shape #!key (requires-grad? #t)) -> tensor</procedure> Creates a 32-bit floating-point tensor. ; data : f32vector holding the elements in row-major order ; shape : list of dimensions, e.g. '(2 3) for a 2x3 matrix or '(10 2 3) for a batch of 10 matrices ; requires-grad? : whether to track gradients (default #t) <enscript highlight="scheme"> ; Single vector (define x (make-tensor32 (f32vector 1.0 2.0 3.0) '(3) requires-grad?: #t)) ; Batch of vectors (define batch (make-tensor32 (make-f32vector 60) '(10 6) requires-grad?: #t)) </enscript> <procedure>(make-tensor64 data shape #!key (requires-grad? #t)) -> tensor</procedure> Creates a 64-bit floating-point tensor from an f64vector. <procedure>(make-base-tensor data shape dtype requires-grad?) -> tensor</procedure> Creates a strict tensor of dtype {{'f32}} or {{'f64}}. This is the constructor behind {{make-tensor32}} and {{make-tensor64}}. ===== Tensor Predicates <procedure>(tensor? obj) -> boolean</procedure> <procedure>(tensor32? obj) -> boolean</procedure> <procedure>(tensor64? obj) -> boolean</procedure> Type predicates for tensors. ===== Tensor Accessors <procedure>(tensor-data tensor) -> vector</procedure> Returns the f32vector or f64vector holding the tensor's data. For a lazy tensor the values are computed if needed. <procedure>(tensor-grad tensor) -> vector or #f</procedure> Returns the gradient vector, or #f if there is no gradient. <procedure>(tensor-shape tensor) -> list</procedure> Returns the shape as a list of dimensions. <procedure>(tensor-dtype tensor) -> symbol</procedure> Returns the data type: {{'f32}} or {{'f64}}. <procedure>(tensor-requires-grad? tensor) -> boolean</procedure> Returns #t if the tensor tracks gradients. <procedure>(tensor-children tensor) -> list</procedure> Returns the input tensors that {{tensor}} was computed from. <procedure>(tensor-backward-fn tensor) -> procedure or #f</procedure> Returns the procedure that sends the gradient of {{tensor}} to its inputs, or #f for a leaf tensor. ===== Arithmetic Operations <procedure>(add a b) -> tensor</procedure> Element-wise addition of tensors a and b. Both tensors must have the same shape and dtype. <enscript highlight="scheme"> (define z (add x y)) ; z = x + y </enscript> Gradient: dL/da = dL/dz, dL/db = dL/dz <procedure>(sub a b) -> tensor</procedure> Element-wise subtraction: a - b. Gradient: dL/da = dL/dz, dL/db = -dL/dz <procedure>(mul a b) -> tensor</procedure> Element-wise multiplication (Hadamard product). Gradient: dL/da = dL/dz * b, dL/db = dL/dz * a <procedure>(div a b) -> tensor</procedure> Element-wise division: a / b. Gradient: dL/da = dL/dz / b, dL/db = -dL/dz * a / b^2 <procedure>(safe-div a b #!key (epsilon 1e-8)) -> tensor</procedure> Element-wise division a / (b + epsilon), which avoids division by zero. ===== Linear Algebra Operations <procedure>(matmul-op a b) -> tensor</procedure> Matrix multiplication using BLAS GEMM/GEMV operations with batch support. Supports: * Matrix x Matrix * Matrix x Vector * Vector x Matrix * Vector x Vector (dot product) * Batched operations (implicit batching over the first dimension) <enscript highlight="scheme"> ; Matrix-vector multiplication (define A (make-tensor32 (f32vector 1.0 2.0 3.0 4.0) '(2 2))) (define b (make-tensor32 (f32vector 5.0 6.0) '(2))) (define c (matmul-op A b)) ; 2x2 matrix times 2-vector = 2-vector ; Batch matrix multiplication (define batch-A (make-tensor32 (make-f32vector 80) '(10 2 4))) ; 10 samples (define W (make-tensor32 (make-f32vector 12) '(4 3))) (define batch-result (matmul-op batch-A W)) ; Shape: (10, 2, 3) </enscript> Gradient: dL/dA = dL/dC . B^T, dL/dB = A^T . dL/dC <procedure>(dot-op a b) -> tensor</procedure> Dot product of two 1-D vectors using BLAS DOT. <enscript highlight="scheme"> (define result (dot-op x y)) ; scalar result </enscript> Gradient: dL/da = dL/dresult * b, dL/db = dL/dresult * a <procedure>(scale-op tensor scalar) -> tensor</procedure> Scalar multiplication using BLAS SCAL. Gradient: dL/dtensor = scalar * dL/dresult ===== Reduction Operations <procedure>(reduce-tensor tensor reducer #!key (compute-gradient #f)) -> tensor</procedure> Generic reduction that keeps gradient flow. The {{reducer}} procedure is applied to each element in the forward pass. The optional {{compute-gradient}} procedure says how the gradient is distributed in the backward pass. ; tensor : input tensor to reduce ; reducer : procedure (element accumulator) -> new-accumulator ; compute-gradient : optional procedure (grad-out index value all-values) -> grad-in. Without it, the gradient is distributed uniformly, as for a sum. Returns a scalar tensor with the reduced value. <enscript highlight="scheme"> ;; Sum all elements (uniform gradient distribution) (define total (reduce-tensor x +)) ;; Product of all elements (gradient uses the product rule) (define prod (reduce-tensor x * compute-gradient: (lambda (grad-out idx val all-values) ;; d(prod)/dx_i = prod / x_i (let ((prod (fold * 1.0 all-values))) (if (> val 0.0) (* grad-out (/ prod val)) 0.0))))) </enscript> <procedure>(sum-tensor tensor) -> tensor</procedure> Sums all elements. The gradient is distributed uniformly to all elements. <procedure>(product-tensor tensor) -> tensor</procedure> Computes the product of all elements. The gradient uses the product rule: d(prod)/dx_i = prod / x_i. <procedure>(mean-tensor tensor) -> tensor</procedure> Computes the mean of all elements. <procedure>(compensated-sum dtype vec start end) -> number</procedure> Sums the elements {{start}} to {{end}} - 1 of the SRFI-4 vector {{vec}} with Kahan compensated summation, which reduces rounding error. ===== Tensor Manipulation Operations <procedure>(slice-tensor tensor start length) -> tensor</procedure> Extracts a slice of a tensor along the first dimension. Gradients flow back to the matching positions of the original tensor. ; tensor : input tensor with shape (n, ...) ; start : starting index (0-based) ; length : number of elements to extract Returns a tensor with shape (length, ...). <enscript highlight="scheme"> ;; Slice a batch of data (define batch-data (make-tensor32 (make-f32vector 100) '(10 10))) (define mini-batch (slice-tensor batch-data 2 5)) ; Shape: (5, 10) ;; Gradients flow back to the original positions (backward! (sum-tensor mini-batch)) (tensor-grad batch-data) ; Only rows 2-6 have non-zero gradients </enscript> <procedure>(reshape tensor new-shape) -> tensor</procedure> Reshapes the tensor. The total number of elements must not change. <procedure>(squeeze tensor #!key (axis #f)) -> tensor</procedure> Removes dimensions of size 1. With {{axis}}, removes only that dimension, which must have size 1; negative values count from the end. If every dimension is removed, the result has shape (1). <procedure>(transpose-tensor tensor axes) -> tensor</procedure> Permutes the dimensions of {{tensor}}. {{axes}} lists the source dimension of each result dimension; for example {{'(0 2 3 1)}} turns (N, C, H, W) into (N, H, W, C). The gradient is sent back through the inverse permutation. <procedure>(flatten-tensor tensor) -> tensor</procedure> Flattens a multi-dimensional tensor to 1-D. ===== Activation Functions <procedure>(relu tensor) -> tensor</procedure> Rectified Linear Unit: max(0, x). Gradient: 1 if x > 0, else 0 <procedure>(tanh-op tensor) -> tensor</procedure> Hyperbolic tangent activation. Gradient: 1 - tanh^2(x) <procedure>(sigmoid tensor) -> tensor</procedure> Sigmoid (logistic) activation: s(x) = 1 / (1 + e^(-x)). Gradient: s(x) * (1 - s(x)) <procedure>(sigmoid-stable tensor) -> tensor</procedure> Sigmoid computed so that it does not overflow for large negative inputs. <procedure>(softmax x #!key (axis -1)) -> tensor</procedure> Numerically stable softmax with batch support. Input shapes: * 1-D: (n_classes) - standard softmax * 2-D: (batch_size, n_classes) - softmax along {{axis}} (default -1, the last axis) <enscript highlight="scheme"> ; Single sample (define logits (make-tensor32 (f32vector 1.0 2.0 3.0) '(3))) (define probs (softmax logits)) ; Sums to 1.0 ; Batch of samples (define batch-logits (make-tensor32 (make-f32vector 60) '(20 3))) (define batch-probs (softmax batch-logits axis: -1)) ; Each row sums to 1.0 </enscript> Gradient: dL/dx = softmax(x) * (dL/dy - sum(dL/dy * softmax(x))) <procedure>(log-softmax x #!key (axis -1)) -> tensor</procedure> Log-softmax with batch support. It is more numerically stable than log(softmax(x)). Input shapes: * 1-D: (n_classes) * 2-D: (batch_size, n_classes) - log-softmax along {{axis}} Gradient: dL/dx = dL/dy - exp(log_softmax(x)) * sum(dL/dy) <procedure>(leaky-relu tensor #!key (alpha 0.01)) -> tensor</procedure> Leaky ReLU: max(alpha * x, x). <procedure>(softplus tensor #!key (beta 1.0)) -> tensor</procedure> Softplus activation: log(1 + e^(beta * x)) / beta. <procedure>(gelu tensor) -> tensor</procedure> Gaussian Error Linear Unit, using the tanh approximation. <procedure>(silu tensor) -> tensor</procedure> SiLU (Sigmoid Linear Unit), also known as Swish: x * s(x). ===== Loss Functions The loss functions scale their gradient by the incoming gradient, so a loss can be scaled or added to other losses before {{backward!}}. <procedure>(mse-loss pred target #!key (reduction 'mean)) -> tensor</procedure> Mean squared error loss with batch support. ; pred : predictions tensor (any shape) ; target : target tensor (same shape as pred) ; reduction : {{'mean}} (average over all elements) or {{'sum}} <enscript highlight="scheme"> ; Single sample (define loss (mse-loss predictions targets)) ; Batch of samples (define batch-pred (make-tensor32 pred-data '(32 10))) (define batch-target (make-tensor32 target-data '(32 10))) (define batch-loss (mse-loss batch-pred batch-target reduction: 'mean)) </enscript> <procedure>(cross-entropy-loss pred target #!key (reduction 'mean) (from-logits #f)) -> tensor</procedure> Cross-entropy loss with batch support. ; pred : predictions: probabilities (softmax already applied) when {{from-logits}} is #f, or raw scores when it is #t ; target : one-hot targets of the same shape as pred, or class indices of shape (batch_size) ; reduction : {{'mean}} (average over the batch) or {{'sum}} ; from-logits : if true, apply log-softmax to pred first Input shapes: * 1-D pred (n_classes): single sample * 2-D pred (batch_size, n_classes): batch of samples <enscript highlight="scheme"> ; Single sample with one-hot target (define loss (cross-entropy-loss probs target)) ; Batch with one-hot targets (define batch-probs (softmax logits axis: -1)) (define batch-loss (cross-entropy-loss batch-probs targets reduction: 'mean)) ; Batch with class indices (define class-indices (make-tensor32 (f32vector 0.0 2.0 1.0) '(3))) (define batch-loss (cross-entropy-loss logits class-indices from-logits: #t reduction: 'mean)) </enscript> ===== Normalization Operations <procedure>(rmsnorm x weight #!key (epsilon 1e-5)) -> tensor</procedure> Root Mean Square layer normalization with batch support. Input shapes: * 1-D: (d_model) - standard RMSNorm * 2-D: (batch_size, d_model) - each row is normalized on its own Formula: output[i] = (x[i] / RMS(x)) * weight[i], where RMS(x) = sqrt(mean(x^2) + epsilon) <enscript highlight="scheme"> ; Single vector (define x (make-tensor32 (make-f32vector 512) '(512))) (define gamma (make-tensor32 (make-f32vector 512 1.0) '(512))) (define normalized (rmsnorm x gamma)) ; Batch of vectors (define batch-x (make-tensor32 (make-f32vector (* 32 512)) '(32 512))) (define batch-norm (rmsnorm batch-x gamma)) ; Normalized per row </enscript> <procedure>(l2-normalize tensor #!key (axis #f) (epsilon 1e-8)) -> tensor</procedure> L2 normalization. ; axis : #f to normalize the entire tensor, or an axis to normalize along For 2-D tensors: * axis 0: each column becomes a unit vector * axis 1: each row becomes a unit vector <enscript highlight="scheme"> ; Normalize the entire tensor (define normalized (l2-normalize x)) ; Normalize each row of a batch (define batch (make-tensor32 (make-f32vector 200) '(10 20))) (define row-normalized (l2-normalize batch axis: 1)) ; Each row has L2 norm 1 </enscript> <procedure>(cosine-similarity a b) -> tensor</procedure> Cosine similarity: (a . b) / (||a|| * ||b||). ===== Convolution Operations <procedure>(conv2d input weight bias #!key (stride 1) (padding 0)) -> tensor</procedure> 2-D convolution using im2col and GEMM, with batch support. ; input : tensor of shape (C_in, H, W) or (N, C_in, H, W) ; weight : tensor of shape (C_out, C_in, KH, KW) ; bias : tensor of shape (C_out), or #f ; stride : convolution stride (default 1) ; padding : zero padding (default 0) Output shapes: * 3-D input: (C_out, H_out, W_out) * 4-D input: (N, C_out, H_out, W_out) <enscript highlight="scheme"> ; Single image (define img (make-tensor32 (make-f32vector (* 3 32 32)) '(3 32 32))) (define output (conv2d img weights bias stride: 2 padding: 1)) ; Batch of images (define batch-imgs (make-tensor32 (make-f32vector (* 16 3 32 32)) '(16 3 32 32))) (define batch-output (conv2d batch-imgs weights bias)) ; Shape: (16, C_out, H_out, W_out) </enscript> ===== Gradient Operations <procedure>(zero-grad! tensor) -> void</procedure> Sets all gradient values to zero. <procedure>(backward! tensor) -> void</procedure> Computes gradients by reverse-mode automatic differentiation. Sorts the computation graph topologically and runs each backward procedure in order. Signals an error if the graph has a cycle. <procedure>(add-to-grad! tensor delta) -> void</procedure> Adds the vector {{delta}} to the tensor's gradient using BLAS AXPY. <procedure>(set-grad! tensor grad) -> void</procedure> Replaces the tensor's gradient vector with {{grad}}. <procedure>(set-backward-fn! tensor fn inputs) -> void</procedure> Makes {{tensor}} the result of a differentiable operation. {{fn}} is a procedure of no arguments that reads {{(tensor-grad tensor)}} and adds the gradient of each input with {{add-to-grad!}}. {{inputs}} is the list of input tensors. This is how new operations are defined: <enscript highlight="scheme"> ;; y = 2x (define (double x) (let ((y (make-base-tensor (f32vector-map (lambda (v) (* 2.0 v)) (tensor-data x)) (tensor-shape x) 'f32 #t))) (set-backward-fn! y (lambda () (add-to-grad! x (f32vector-map (lambda (g) (* 2.0 g)) (tensor-grad y)))) (list x)) y)) </enscript> {{f32vector-map}} in this example stands for any element-wise map over an f32vector; it is not part of nanograd. ===== Utility Functions <procedure>(tensor->list tensor) -> list</procedure> Converts tensor data to a list. <procedure>(print-tensor tensor) -> void</procedure> Prints the tensor's shape, dtype, data and gradient. <procedure>(vector-length-for-dtype vec dtype) -> integer</procedure> Returns the length of an f32vector ({{'f32}}) or f64vector ({{'f64}}). The module also exports small helpers used by the other nanograd modules: {{f32vector-fold}}, {{f64vector-fold}}, and the BLAS-style vector routines {{ssub}}, {{dsub}}, {{sblit}}, {{dblit}}, {{scopy}} and {{dcopy}}. ==== nanograd-layer Neural network layers and containers with batch processing support. ===== Layer Predicates <procedure>(layer? obj) -> boolean</procedure> <procedure>(dense-layer? obj) -> boolean</procedure> <procedure>(conv2d-layer? obj) -> boolean</procedure> <procedure>(batch-norm-2d? obj) -> boolean</procedure> <procedure>(sequential? obj) -> boolean</procedure> <procedure>(flatten-layer? obj) -> boolean</procedure> ===== Dense Layer <procedure>(make-dense-layer input-size output-size #!key (activation (make-identity)) (use-bias #t) (dtype 'f32) (name "Dense")) -> layer</procedure> Creates a fully connected layer. Weights start as uniform random values scaled by sqrt(2 / input-size); biases start at zero. ; input-size : number of input features ; output-size : number of output features ; activation : activation object (default identity) ; use-bias : whether to include a bias term (default #t) ; dtype : {{'f32}} or {{'f64}} (default {{'f32}}) ; name : layer name Input shapes: * 1-D: (input_size) -> output: (output_size) * 2-D: (batch_size, input_size) -> output: (batch_size, output_size) For 2-D inputs, the layer uses one BLAS GEMM for the whole batch. <enscript highlight="scheme"> (define layer (make-dense-layer 784 128 activation: (make-relu) name: "Hidden1")) ; Single input (define x (make-tensor32 (make-f32vector 784) '(784))) (define output (forward layer x)) ; Shape: (128) ; Batch input (define batch-x (make-tensor32 (make-f32vector (* 32 784)) '(32 784))) (define batch-output (forward layer batch-x)) ; Shape: (32, 128) </enscript> ===== Convolutional Layer <procedure>(make-conv2d-layer in-channels out-channels kernel-size #!key (stride 1) (padding 0) (activation (make-identity)) (dtype 'f32) (weight-values #f) (bias-values #f) (name "Conv2D")) -> layer</procedure> Creates a 2-D convolutional layer. Unless {{weight-values}} is given, weights start as uniform random values scaled by sqrt(2 / fan-in) (He scaling). ; in-channels : number of input channels ; out-channels : number of output channels ; kernel-size : size of the square kernel ; stride : convolution stride (default 1) ; padding : zero padding (default 0) ; activation : activation object ; dtype : {{'f32}} or {{'f64}} ; weight-values : optional SRFI-4 vector of initial weights, laid out as (out-channels, in-channels, kernel-size, kernel-size) ; bias-values : optional SRFI-4 vector of initial biases ; name : layer name Input shapes: * 3-D: (C_in, H, W) - single image * 4-D: (N, C_in, H, W) - batch of images Output shapes: * 3-D: (C_out, H_out, W_out) * 4-D: (N, C_out, H_out, W_out) <enscript highlight="scheme"> (define conv (make-conv2d-layer 3 32 3 stride: 1 padding: 1 activation: (make-relu))) ; Single image (define img (make-tensor32 img-data '(3 32 32))) (define features (forward conv img)) ; Shape: (32, 32, 32) ; Batch of images (define batch (make-tensor32 batch-data '(16 3 32 32))) (define batch-features (forward conv batch)) ; Shape: (16, 32, 32, 32) </enscript> ===== Batch Normalization Layer <procedure>(make-batch-norm-2d num-features #!key (epsilon 1e-5) (momentum 0.1) (dtype 'f32) (name "BatchNorm2d")) -> layer</procedure> Creates a 2-D batch normalization layer: y = gamma * (x - mu) / sqrt(sigma^2 + epsilon) + beta where mu and sigma^2 are the per-channel mean and variance of the batch in training mode, and running averages of them in evaluation mode. ; num-features : number of channels (C) ; epsilon : small constant for numerical stability (default 1e-5) ; momentum : weight of the current batch in the running statistics (default 0.1) ; dtype : {{'f32}} or {{'f64}} (default {{'f32}}) ; name : layer name Input shapes: * 3-D: (C, H, W) - treated as a batch of 1 * 4-D: (N, C, H, W) - standard batch The output has the same shape as the input. gamma and beta are learned. <enscript highlight="scheme"> ;; Batch norm for 64 channels (define bn (make-batch-norm-2d 64 epsilon: 1e-5 momentum: 0.1)) ;; Training mode: uses batch statistics (set-training-mode! bn #t) (define normalized (forward bn input)) ; Input shape: (N, 64, H, W) ;; Evaluation mode: uses running statistics (set-eval-mode! bn) (define test-normalized (forward bn test-input)) </enscript> ===== Layer Normalization <procedure>(layer-norm x gamma beta #!key (epsilon 1e-5)) -> tensor</procedure> Normalizes each row of {{x}} over its last axis. Each row x_r of length D is mapped to y_r = gamma * (x_r - mean(x_r)) / sqrt(var(x_r) + epsilon) + beta where var is the biased variance. {{gamma}} and {{beta}} are tensors of shape (D). Gradients flow to {{x}}, {{gamma}} and {{beta}}. ===== Flatten Layer <procedure>(make-flatten #!key (name "Flatten")) -> layer</procedure> Creates a layer that flattens each sample to a vector. * 4-D: (N, C, H, W) -> (N, C*H*W) * 3-D: (C, H, W) -> (C*H*W) * 2-D and 1-D inputs are passed through unchanged <enscript highlight="scheme"> (define flatten (make-flatten name: "Flatten")) ; Flatten a batch of feature maps (define features (make-tensor32 data '(32 64 8 8))) (define flattened (forward flatten features)) ; Shape: (32, 4096) </enscript> ===== Pooling <procedure>(maxpool2d input kernel-size #!key (stride #f)) -> tensor</procedure> 2-D max pooling with a square window of size {{kernel-size}}. The stride defaults to the kernel size, so windows do not overlap. * 3-D: (C, H, W) -> (C, OH, OW) * 4-D: (N, C, H, W) -> (N, C, OH, OW) The gradient goes to the largest element of each window. <procedure>(global-avg-pool2d input) -> tensor</procedure> Averages over the spatial dimensions. * 3-D: (C, H, W) -> (C) * 4-D: (N, C, H, W) -> (N, C) The gradient is distributed uniformly over the spatial positions of each channel. <enscript highlight="scheme"> ;; Batch of feature maps (define batch-features (make-tensor32 (make-f32vector (* 32 128 8 8)) '(32 128 8 8))) (define batch-pooled (global-avg-pool2d batch-features)) ; Shape: (32, 128) ;; Classification head (define logits (forward fc-layer batch-pooled)) ; Shape: (32, num_classes) </enscript> Global average pooling can replace a large fully connected layer at the end of a CNN, with far fewer parameters. ===== Sequential Container <procedure>(make-sequential layers #!key (name "Sequential")) -> layer</procedure> Creates a container that applies {{layers}} in order. <enscript highlight="scheme"> (define model (make-sequential (list (make-dense-layer 784 128 activation: (make-relu)) (make-dense-layer 128 64 activation: (make-relu)) (make-dense-layer 64 10 activation: (make-identity))) name: "MLP")) ; Works with both single and batch inputs (define single-output (forward model single-input)) (define batch-output (forward model batch-input)) </enscript> ===== Layer Operations <procedure>(forward layer input) -> tensor</procedure> Runs the layer on {{input}}, which may be a single sample or a batch. <procedure>(parameters layer) -> list</procedure> Returns the list of trainable parameter tensors. <procedure>(zero-grad-layer! layer) -> void</procedure> Zeros the gradients of all parameters of the layer. <procedure>(set-training-mode! layer training?) -> void</procedure> Puts the layer, and every layer inside a container, in training mode ({{training?}} is #t) or evaluation mode (#f). In training mode batch normalization uses batch statistics and updates its running statistics; in evaluation mode it uses the running statistics and leaves them unchanged. <procedure>(set-eval-mode! layer) -> void</procedure> Same as {{(set-training-mode! layer #f)}}. <procedure>(layer-input-size layer) -> integer or #f</procedure> <procedure>(layer-output-size layer) -> integer or #f</procedure> <procedure>(layer-activation layer) -> activation</procedure> <procedure>(layer-name layer) -> string</procedure> Layer properties. The sizes are #f for layers whose sizes depend on the input, such as flatten. ===== Activation Objects <procedure>(make-relu) -> activation</procedure> <procedure>(make-tanh) -> activation</procedure> <procedure>(make-sigmoid) -> activation</procedure> <procedure>(make-gelu) -> activation</procedure> <procedure>(make-silu) -> activation</procedure> <procedure>(make-identity) -> activation</procedure> Create activation objects for use in layers. <procedure>(activation? obj) -> boolean</procedure> <procedure>(activation-forward act x) -> tensor</procedure> <procedure>(activation-name act) -> string</procedure> ===== Checkpoints Layers and models are saved to and loaded from CBOR files. <procedure>(save-model model filepath) -> void</procedure> <procedure>(save-layer layer filepath) -> void</procedure> Write {{model}} (a sequential container or a single layer) with all its weights to {{filepath}}. The two procedures are the same. <procedure>(load-model filepath) -> layer</procedure> <procedure>(load-layer filepath) -> layer</procedure> Read a model written by {{save-model}}. They signal an error if the file is not a nanograd checkpoint or has an unsupported version. <procedure>(layer->serializable layer) -> alist</procedure> <procedure>(serializable->layer alist) -> layer</procedure> Convert a layer to and from the association list that is stored in the file. Dense, convolutional, flatten and sequential layers can be saved. The AM-backed layers of {{nanograd-array-morphisms}} cannot. <enscript highlight="scheme"> (save-model model "model.ngrd") (define restored (load-model "model.ngrd")) </enscript> A checkpoint file holds one CBOR data item (RFC 8949), written with the [[cbor]] egg, after the CBOR self-describe tag. The item is an association list: <enscript highlight="scheme"> ((format . nanograd-checkpoint) (version . 1) (model . <output of layer->serializable>)) </enscript> Tensor data are stored as RFC 8746 typed arrays (tag 85 for little-endian f32, tag 86 for little-endian f64), so they are read and written as single blocks of bytes and round-trip bit for bit. Other CBOR decoders can read the files; for example, Python's {{cbor2}} package decodes a checkpoint, and {{numpy.frombuffer}} turns each typed array into an array. Checkpoints written by nanograd 3.0 and earlier (CHICKEN 5, s11n format) cannot be read. ===== Utility Functions <procedure>(print-layer layer #!optional (indent 0)) -> void</procedure> Prints layer information with optional indentation. <procedure>(summary model) -> void</procedure> Prints a model summary with its layers and parameters. ==== nanograd-optimizer Optimization algorithms for neural network training. ===== Optimizer Predicates <procedure>(optimizer? obj) -> boolean</procedure> <procedure>(sgd? obj) -> boolean</procedure> <procedure>(adam? obj) -> boolean</procedure> <procedure>(rmsprop? obj) -> boolean</procedure> ===== SGD Optimizer <procedure>(make-sgd parameters #!key (learning-rate 0.01) (momentum 0.0) (weight-decay 0.0) (nesterov #f)) -> optimizer</procedure> Stochastic gradient descent with optional momentum and Nesterov momentum. ; parameters : list of parameter tensors to optimize ; learning-rate : step size (default 0.01) ; momentum : momentum factor (default 0.0, no momentum) ; weight-decay : L2 regularization factor (default 0.0) ; nesterov : use Nesterov momentum (default #f) ===== Adam Optimizer <procedure>(make-adam parameters #!key (learning-rate 0.001) (beta1 0.9) (beta2 0.999) (epsilon 1e-8) (weight-decay 0.0)) -> optimizer</procedure> Adam (adaptive moment estimation) with bias correction. ; beta1 : decay rate of the first moment (default 0.9) ; beta2 : decay rate of the second moment (default 0.999) ; epsilon : numerical stability constant (default 1e-8) ; weight-decay : L2 regularization factor, added to the gradient (default 0.0) ===== RMSprop Optimizer <procedure>(make-rmsprop parameters #!key (learning-rate 0.01) (alpha 0.99) (epsilon 1e-8) (weight-decay 0.0) (momentum 0.0)) -> optimizer</procedure> RMSprop with optional momentum. ; alpha : smoothing constant (default 0.99) ===== Optimizer Operations <procedure>(step! optimizer) -> void</procedure> Updates the parameters from their current gradients. <procedure>(get-learning-rate optimizer) -> number</procedure> Returns the current learning rate. <procedure>(set-learning-rate! optimizer lr) -> void</procedure> Sets the learning rate, for example from a learning rate schedule. <procedure>(optimizer-state optimizer) -> alist</procedure> Returns the optimizer's settings as an association list (for Adam, also the step count). ===== Optimizer Kernels An optimizer kernel performs an optimizer's whole update of one parameter vector in a single call, for example in compiled C code. When a kernel is registered, {{step!}} calls it instead of the Scheme update. Currently Adam uses registered kernels. <procedure>(register-optimizer-kernel! name dtype kernel) -> void</procedure> Uses {{kernel}} for the updates of optimizer {{name}} (a symbol, such as {{'adam}}) on parameters of dtype {{'f32}} or {{'f64}}. A kernel of #f removes the registration. An Adam kernel is a procedure {{(kernel n hyper param grad m v)}} that updates the first {{n}} elements of the SRFI-4 vectors {{param}}, {{m}} and {{v}} in place, and adds the weight-decay term to {{grad}}. {{hyper}} is an f64vector holding the learning rate, beta1, beta2, epsilon, the bias corrections 1 - beta1^t and 1 - beta2^t, and the weight decay. <procedure>(optimizer-kernel name dtype) -> procedure or #f</procedure> Returns the kernel registered for optimizer {{name}} on {{dtype}} parameters, or #f. ==== nanograd-training Learning rate schedules and gradient clipping. ===== Learning Rate Schedules A schedule is a procedure that takes a step number and returns the learning rate for that step. <procedure>(make-constant-schedule learning-rate) -> procedure</procedure> Always returns {{learning-rate}}. <procedure>(make-linear-warmup-schedule init-lr final-lr warmup-steps) -> procedure</procedure> Rises linearly from {{init-lr}} to {{final-lr}} over {{warmup-steps}} steps, then stays at {{final-lr}}. <procedure>(make-cosine-schedule max-lr min-lr total-steps #!key (warmup-steps 0)) -> procedure</procedure> Cosine annealing from {{max-lr}} down to {{min-lr}} over {{total-steps}}, after an optional linear warmup: lr = min-lr + 0.5 * (max-lr - min-lr) * (1 + cos(pi * progress)), where progress = (step - warmup-steps) / (total-steps - warmup-steps) <procedure>(make-cosine-warmup-schedule max-lr min-lr total-steps warmup-steps) -> procedure</procedure> Same as {{(make-cosine-schedule max-lr min-lr total-steps warmup-steps: warmup-steps)}}. <procedure>(make-step-decay-schedule init-lr decay-rate decay-steps #!key (warmup-steps 0)) -> procedure</procedure> Multiplies the learning rate by {{decay-rate}} every {{decay-steps}} steps, after an optional warmup: lr = init-lr * decay-rate ^ floor(step / decay-steps) <procedure>(learning-rate-at-step schedule step) -> number</procedure> Same as {{(schedule step)}}. <enscript highlight="scheme"> (define schedule (make-cosine-schedule 1e-3 1e-5 10000 warmup-steps: 500)) (set-learning-rate! optimizer (schedule step)) </enscript> ===== Gradient Clipping <procedure>(compute-gradient-norm parameters) -> number</procedure> Returns the global L2 norm of the gradients of {{parameters}}: the square root of the sum of all squared gradient elements. <procedure>(clip-gradients-by-norm! parameters max-norm) -> number</procedure> If the global L2 norm is larger than {{max-norm}}, scales all gradients so that it equals {{max-norm}}, which keeps their direction. Returns the norm before clipping. <procedure>(clip-gradients-by-value! parameters clip-value) -> void</procedure> Clamps each gradient element to the range [-clip-value, clip-value]. <enscript highlight="scheme"> (backward! loss) (clip-gradients-by-norm! (parameters model) 1.0) (step! optimizer) </enscript> ==== nanograd-diagnostics Gradient statistics, health checks and monitoring over training. Both this module and {{nanograd-training}} export {{compute-gradient-norm}}. To import both modules, rename or exclude one of them, for example: <enscript highlight="scheme"> (import nanograd-training (except nanograd-diagnostics compute-gradient-norm)) </enscript> ===== Gradient Statistics <procedure>(gradient-stats parameters) -> alist</procedure> Returns an association list with the global L2 norm ({{global-l2-norm}}), the global infinity norm ({{global-inf-norm}}), and {{layer-stats}}: for each parameter with a gradient, its index ({{layer}}), {{l2-norm}}, {{inf-norm}}, mean absolute value ({{mean-abs}}) and {{size}}. <procedure>(compute-gradient-norm parameters #!key (norm-type 'l2)) -> number</procedure> Returns the global gradient norm: L2 ({{'l2}}) or infinity ({{'inf}}). <procedure>(gradient-l2-norm grad dtype) -> number</procedure> <procedure>(gradient-inf-norm grad dtype) -> number</procedure> Return the L2 norm or the largest absolute value of one gradient vector. ===== Health Checks <procedure>(has-exploding-gradients? parameters #!key (threshold 10.0) (norm-type 'l2)) -> boolean</procedure> Returns #t if the global gradient norm is above {{threshold}}, or is NaN or infinite. <procedure>(has-vanishing-gradients? parameters #!key (threshold 1e-7) (norm-type 'l2)) -> boolean</procedure> Returns #t if the global gradient norm is below {{threshold}}. <procedure>(check-gradient-health parameters #!key (exploding-threshold 10.0) (vanishing-threshold 1e-7) (imbalance-threshold 1000.0) (warn? #t)) -> alist</procedure> Checks for exploding, vanishing and unbalanced gradients. Gradients are unbalanced when the largest per-parameter L2 norm is more than {{imbalance-threshold}} times the smallest one; this check is skipped when some gradient is zero. With {{warn?}}, each problem found is printed. Returns an association list with {{healthy?}} (#t if no problem was found), {{issues}} (a list of symbols, such as {{gradient-imbalance}}) and {{stats}} (the result of {{gradient-stats}}). <procedure>(print-gradient-report parameters #!optional (label "")) -> void</procedure> Prints the health status and gradient statistics. ===== Gradient Monitor <procedure>(make-gradient-monitor #!key (exploding-threshold 10.0) (vanishing-threshold 1e-7) (imbalance-threshold 1000.0) (history-size 100)) -> monitor</procedure> Creates an object that records gradient statistics over training. It keeps the last {{history-size}} steps. <procedure>(gradient-monitor? obj) -> boolean</procedure> <procedure>(record-step! monitor step parameters) -> alist</procedure> Records the gradient norms and health of {{parameters}} at training step {{step}}. Returns the health check result, as from {{check-gradient-health}}. <procedure>(get-gradient-history monitor) -> list</procedure> Returns the recorded steps, oldest first. Each entry is an association list with {{step}}, {{global-l2-norm}}, {{global-inf-norm}}, {{healthy?}} and {{issues}}. <procedure>(diagnose-training monitor) -> alist</procedure> Summarizes the last 10 recorded steps: mean, largest and smallest gradient norm, the number of unhealthy steps and warnings, a stability ratio (smallest norm divided by largest), and a list of {{recommendations}} as strings. ===== Remediation <procedure>(clip-gradients! parameters #!key (max-norm 1.0) (norm-type 'l2) (verbose? #t)) -> number or #f</procedure> Scales the gradients down if their global norm is above {{max-norm}}. Returns the norm before clipping if clipping was applied, and #f otherwise. <procedure>(suggest-learning-rate parameters current-lr gradient-norm) -> number</procedure> Suggests a new learning rate from the gradient norm: lower for large norms, higher for very small ones. Prints the suggestion. <enscript highlight="scheme"> (import (except nanograd-diagnostics compute-gradient-norm)) (define monitor (make-gradient-monitor)) ;; In the training loop, after backward! (record-step! monitor step (parameters model)) ;; After training (diagnose-training monitor) </enscript> ==== nanograd-array-morphisms Lazy tensors backed by [[array-morphisms]], layers built on them, and training and inference steps that are compiled once and replayed. A lazy tensor wraps an array-morphisms variable. Operations on lazy tensors build an expression instead of computing values; the values are computed when they are read with {{tensor-data}}, or when a step procedure below runs the computation. Because array-morphisms sees the whole computation, it can fuse operations, reuse buffers, and run matrix products and element-wise operations on registered C backends. The module provides three ways to train: ; {{am-training-step/ssa}} : Compiles the forward pass, the loss and the backward pass of the model into one SSA program on the first call, and replays it on later calls. This is the fastest way. ; {{am-training-step}} : Runs the forward and backward passes layer by layer through two morphism contexts, which reuse buffers after the first step. ; {{backward!}} on a lazy loss : Builds the gradients as lazy expressions, without contexts. For inference, {{am-forward/ssa}} compiles and replays the forward pass alone. Morphism contexts are created with {{make-morphism-context}} from the {{array-morphisms-context}} module. Each context holds one compiled program, and a compiled program is fixed to the shapes of the inputs it was compiled for. Use one context per model and input shape, for example one per batch size. ===== Lazy Tensors <procedure>(make-lazy-tensor mv) -> tensor</procedure> Creates a lazy tensor wrapping the array-morphisms variable {{mv}}. <procedure>(get-or-make-lazy mv) -> tensor</procedure> Returns the lazy tensor for {{mv}}, creating it if needed. Use this procedure to wrap input data: <enscript highlight="scheme"> (import array-morphisms-core (prefix array-morphisms-grad am:)) ;; A batch of 32 grayscale 28x28 images, without copying the data (define x-lt (get-or-make-lazy (am:make-var (make-morphism image-data '(32 1 28 28) 'f32) #f))) ; no gradient for the input </enscript> <procedure>(lazy-tensor-morph-variable tensor) -> morph-variable or #f</procedure> Returns the array-morphisms variable of a lazy tensor, or #f if {{tensor}} is not lazy. <procedure>(activation-forward-am act tensor) -> tensor</procedure> Applies the activation object {{act}} (from {{make-relu}}, {{make-gelu}}, and so on) to a lazy tensor. ===== Loss Functions <procedure>(am-mse-loss pred target) -> tensor</procedure> Mean squared error of two lazy tensors. Returns a scalar lazy tensor. <procedure>(am-cross-entropy-loss logits target) -> tensor</procedure> Numerically stable softmax cross-entropy, averaged over the batch. {{logits}} is a (batch, C) lazy tensor of raw scores and {{target}} a (batch, C) lazy tensor of one-hot labels. Returns a scalar lazy tensor. ===== AM-backed Layers These layers take and return lazy tensors, and answer the same layer operations as the layers of {{nanograd-layer}} ({{forward}}, {{parameters}}, {{set-training-mode!}}, and so on). Their parameters are lazy tensors, so the optimizers of {{nanograd-optimizer}} work with them. Their weights start from He initialization. They cannot be saved with {{save-model}}. <procedure>(make-am-dense-layer in-size out-size #!key (activation (make-identity)) (dtype 'f64)) -> layer</procedure> Fully connected layer for inputs of shape (batch, in-size). <procedure>(make-am-conv2d-layer in-channels out-channels kernel-size #!key (stride 1) (padding 0) (activation (make-identity)) (dtype 'f64)) -> layer</procedure> 2-D convolution, run as a single operation that the SSA compiler differentiates. {{kernel-size}} is a number, or a pair (KH . KW). The input has shape (N, C, H, W), or (N, H, W, C) when its last dimension equals {{in-channels}}. '''The output has shape (N, OH, OW, C_out)''', with channels last, so that the next convolution can read it without a copy. A flatten layer after it therefore orders the features as (H, W, C), not (C, H, W) as for {{make-conv2d-layer}}. <procedure>(make-am-batch-norm num-features #!key (epsilon 1e-5) (momentum 0.1) (dtype 'f64)) -> layer</procedure> Batch normalization over the batch axis (axis 0), with learned gamma and beta of length {{num-features}}. <procedure>(make-am-flatten #!key (name "Flatten")) -> layer</procedure> Flattens each sample: (N, ...) -> (N, product of the other dimensions). 1-D and 2-D inputs are passed through unchanged. <procedure>(make-am-sequential layers) -> layer</procedure> Applies {{layers}} in order. <procedure>(am-parameters layer) -> list</procedure> Returns the parameters of {{layer}} as lazy tensors. <procedure>(am-zero-grad! layer) -> void</procedure> Zeros the gradients of all parameters of {{layer}}. ===== Compiled Training and Inference <procedure>(am-training-step/ssa ctx optimizer model loss-fn x . extra-inputs) -> tensor</procedure> Performs one training step: runs {{model}} on the lazy tensor {{x}}, computes the loss with {{(loss-fn output)}}, computes the gradients, calls {{step!}} on {{optimizer}} and zeros the gradients. Returns the loss as a lazy tensor. The first call with a context {{ctx}} compiles the forward pass, the loss and the backward pass into one SSA program and plans its buffers; later calls replay that program with the new input. Parameters updated by the optimizer are seen by the next replay automatically. {{extra-inputs}} are further lazy tensors that change every step, such as the targets. Pass each one here and also use it inside {{loss-fn}}, so that the compiled program reads its new value on every step. The procedure signals an error if some parameter of the model gets no gradient. <procedure>(am-forward/ssa ctx model x) -> tensor</procedure> Runs {{model}} on the lazy tensor {{x}} for inference, through a compiled forward-only program. The first call with {{ctx}} compiles the program; later calls replay it with the new input. The program reads the model's current weights on every call. Use one context per input shape, and do not use a context for both {{am-forward/ssa}} and {{am-training-step/ssa}}. The result's data lives in the context's buffers and is overwritten by the next call with the same context, so read it with {{tensor-data}} before calling again. <procedure>(am-replay-plan-stats ctx) -> alist</procedure> Returns an association list of (instruction-tag . count) for the replay plan compiled for {{ctx}}. ===== Layer-by-layer Training with Contexts <procedure>(make-am-training-context) -> (values ctx-fwd ctx-bwd)</procedure> Returns two new morphism contexts, for the forward and the backward pass. <procedure>(am-training-step ctx-fwd ctx-bwd optimizer model loss-fn x) -> tensor</procedure> Performs one training step through the two contexts: the forward pass layer by layer, the loss, the backward pass, the optimizer step, and zeroing of the gradients. After the first step both contexts reuse the buffers they planned. Returns the loss as a lazy tensor. <procedure>(am-forward/ctx ctx layer x) -> tensor</procedure> Runs {{layer}} on {{x}} through {{ctx}}, computing each layer's output before the next layer runs. <procedure>(am-backward!/ctx ctx loss) -> void</procedure> Runs the backward pass from the lazy scalar {{loss}} through {{ctx}}, computing each gradient before it is passed on. <procedure>(am-realize-grads/ctx ctx parameters) -> void</procedure> Computes the gradients of all {{parameters}} through {{ctx}}, so that buffers can be shared among them. <procedure>(am-finalize-context! ctx) -> void</procedure> Plans the buffers of {{ctx}} from the allocations recorded so far. <procedure>(am-reset-context! ctx) -> void</procedure> Prepares {{ctx}} for the next step, which reuses the planned buffers. <procedure>(fresh-copy-morphism m) -> morphism</procedure> Copies the concrete array {{m}} into a new row-major buffer that does not belong to any context. ===== Kernel Backends array-morphisms can run matrix products, convolutions and element-wise operations on registered C backends; see {{register-blas-backend!}} and {{register-activation-backend!}} in [[array-morphisms]]. nanograd uses whatever backends are registered. Together with an Adam kernel registered with {{register-optimizer-kernel!}}, this moves nearly all of the work of a compiled training step out of Scheme. === Examples ==== Batch Processing with Dense Layers <enscript highlight="scheme"> (import nanograd-autograd nanograd-layer) (define batch-size 32) (define input-dim 784) (define batch-data (make-f32vector (* batch-size input-dim))) ;; Fill with data... (define batch-input (make-tensor32 batch-data (list batch-size input-dim))) (define layer (make-dense-layer input-dim 128 activation: (make-relu))) (define output (forward layer batch-input)) ; Shape: (32, 128) </enscript> ==== Training an MLP <enscript highlight="scheme"> (import nanograd-autograd nanograd-layer nanograd-optimizer) (define model (make-sequential (list (make-dense-layer 784 256 activation: (make-relu)) (make-dense-layer 256 128 activation: (make-relu)) (make-dense-layer 128 10 activation: (make-identity))) name: "MLP")) (define optimizer (make-adam (parameters model) learning-rate: 0.001)) ;; train-batches is a list of (x . y) pairs: ;; x of shape (batch_size, 784), y one-hot of shape (batch_size, 10) (define (train-epoch train-batches) (set-training-mode! model #t) (for-each (lambda (batch) (let* ((logits (forward model (car batch))) (loss (cross-entropy-loss logits (cdr batch) from-logits: #t reduction: 'mean))) (backward! loss) (step! optimizer) (zero-grad-layer! model))) train-batches)) </enscript> ==== Convolutional Network with Batch Normalization <enscript highlight="scheme"> (import nanograd-autograd nanograd-layer nanograd-optimizer) (define cnn (make-sequential (list (make-conv2d-layer 3 32 3 stride: 1 padding: 1) (make-batch-norm-2d 32) (make-conv2d-layer 32 64 3 stride: 1 padding: 1) (make-batch-norm-2d 64) (make-flatten) (make-dense-layer (* 64 32 32) 128 activation: (make-relu)) (make-dense-layer 128 10)) name: "CNN")) (define batch-images (make-tensor32 image-data '(16 3 32 32))) ; 16 RGB images (set-training-mode! cnn #t) (define predictions (forward cnn batch-images)) ; Shape: (16, 10) </enscript> ==== Saving and Restoring a Model <enscript highlight="scheme"> (import nanograd-layer) (save-model cnn "cnn.ngrd") (define restored (load-model "cnn.ngrd")) (set-eval-mode! restored) </enscript> ==== Compiled Training of a CNN <enscript highlight="scheme"> (import srfi-4 array-morphisms-core array-morphisms-context (prefix array-morphisms-grad am:) nanograd-autograd nanograd-layer nanograd-optimizer nanograd-array-morphisms) (define num-classes 4) (define model (make-am-sequential (list (make-am-conv2d-layer 1 16 3 padding: 1 activation: (make-relu) dtype: 'f32) (make-am-conv2d-layer 16 32 3 stride: 2 padding: 1 activation: (make-relu) dtype: 'f32) (make-am-flatten) (make-am-dense-layer (* 14 14 32) 64 activation: (make-relu) dtype: 'f32) (make-am-dense-layer 64 num-classes dtype: 'f32)))) (define optimizer (make-adam (am-parameters model) learning-rate: 0.001)) ;; Wraps an f32vector as a lazy input tensor, without copying (define (lazy-input data shape) (get-or-make-lazy (am:make-var (make-morphism data shape 'f32) #f))) ;; One context for training at batch size 32 (define train-ctx (make-morphism-context)) ;; images: f32vector of 32*1*28*28 values; labels: f32vector of 32*4 one-hot values (define (train-batch images labels) (let* ((x (lazy-input images '(32 1 28 28))) (y (lazy-input labels (list 32 num-classes))) (loss (am-training-step/ssa train-ctx optimizer model (lambda (logits) (am-cross-entropy-loss logits y)) x y))) (f32vector-ref (tensor-data loss) 0))) ;; A separate context for inference at batch size 32 (define eval-ctx (make-morphism-context)) (define (predict images) (tensor-data (am-forward/ssa eval-ctx model (lazy-input images '(32 1 28 28))))) </enscript> === Performance Notes * The strict tensors of {{nanograd-autograd}} use BLAS for matrix operations, including batched GEMM. Batched operations are much faster than processing samples one at a time. * {{am-training-step/ssa}} compiles a training step once and replays it with planned buffers, which avoids most allocation and runs fused operations. With C kernel backends registered in array-morphisms and an Adam kernel registered with {{register-optimizer-kernel!}}, almost no per-element work is left in Scheme. * Use f32 tensors when 64-bit precision is not needed. === Limitations * CPU only (no GPU support). * Batches must be built by the program. * The strict tensor operations run in a single thread; only the registered array-morphisms backends can use several threads. * AM-backed layers cannot be saved to checkpoints. * A compiled program is fixed to the input shape it was compiled for; each input shape needs its own morphism context. * Batch normalization needs explicit switching between training and evaluation mode. === Troubleshooting '''Shape mismatch errors''' Make sure tensor shapes match. For batched operations, the batch dimensions must agree. <enscript highlight="scheme"> (define x (make-tensor32 (make-f32vector 200) '(10 20))) (define y (make-tensor32 (make-f32vector 300) '(15 20))) (add x y) ; Error: shape mismatch </enscript> '''A compiled step rejects a batch''' {{am-forward/ssa}} signals an error when the input shape differs from the shape the context was compiled for, as happens with a smaller last batch. Use a separate context for each batch size. '''Batch normalization mode not set''' Set the mode before training and before evaluation: <enscript highlight="scheme"> (set-training-mode! model #t) (train-epoch model) (set-eval-mode! model) (evaluate model) </enscript> === Author [[https://github.com/iraikov|Ivan Raikov]] === Repository [[https://github.com/iraikov/nanograd|https://github.com/iraikov/nanograd]] === Version History ; 3.1 : Support for CHICKEN 6 * Checkpoints use the CBOR format of the [[cbor]] egg, replacing s11n; older checkpoints cannot be read * Optimizer kernels: {{register-optimizer-kernel!}} and {{optimizer-kernel}}; Adam uses a registered kernel * {{am-forward/ssa}}: compiled forward pass for inference * Batched dense layers add the bias to every row and compute the bias gradient from every row * {{transpose-tensor}} has a backward pass * Corrected gradient of {{rmsnorm}} with respect to its input * {{mse-loss}} and {{cross-entropy-loss}} scale their gradient by the incoming gradient * {{layer-norm}} normalizes each row over the last axis, with gradients to x, gamma and beta * Corrected lazy GeLU; {{am-training-step/ssa}} signals an error when a parameter gets no gradient * Corrections in nanograd-diagnostics ; 3.0 : array-morphisms integration * Module nanograd-array-morphisms: lazy tensors, AM-backed layers, context-based buffer reuse * {{am-training-step/ssa}}: training steps compiled to an SSA program and replayed * Differentiable conv2d as a single SSA operation ; 2.3 : Iris dataset example ; 2.2 : Module nanograd-training: learning rate schedules and gradient clipping ; 2.1 : Unit tests use the test egg * Gradient computation in the batch norm layer * Batched maxpool2d ; 2.0 : Batch processing support * Dense layers support 1-D and 2-D inputs * Conv2D supports 3-D and 4-D inputs * Batch normalization for 3-D and 4-D inputs * Softmax and log-softmax with batch and axis support * Cross-entropy loss with batch reduction * RMSNorm with 1-D and 2-D support * Global average pooling with 3-D and 4-D support * L2-normalize with axis parameter ; 1.2 : Additional operations * Reduction operations (sum-tensor, mean-tensor, product-tensor, reduce-tensor) * Tensor slicing (slice-tensor) * Batch normalization (make-batch-norm-2d) * Global average pooling (global-avg-pool2d) * Training and evaluation mode control ; 1.1 : Correction in the mul operation ; 1.0 : Initial release * Core autograd engine * Dense and convolutional layers * SGD, Adam and RMSprop optimizers * Basic activation and loss functions === See Also * [[array-morphisms]] * [[cbor]] * [[blas]] * [[yasos]] * [[mathh]] === License LGPL-3 === References * PyTorch: dynamic computation graphs, autograd design, and batch-first conventions * micrograd: minimalist autograd engine by Andrej Karpathy * "Automatic Differentiation in Machine Learning: a Survey" (Baydin et al., 2018) * "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift" (Ioffe and Szegedy, 2015) * RFC 8949: Concise Binary Object Representation (CBOR) * RFC 8746: Concise Binary Object Representation (CBOR) Tags for Typed Arrays
Description of your changes:
I would like to authenticate
Authentication
Username:
Password:
Spam control
What do you get when you subtract 4 from 18?